Skip to main content

krishiv_sql/catalog/
mod.rs

1//! Catalog abstractions for Krishiv.
2//!
3//! This crate defines `TableProvider`, `CatalogProvider`, schema types, and
4//! column statistics. An in-memory reference implementation is included.
5
6pub mod iceberg_rest;
7pub use iceberg_rest::{
8    GenericRestCatalog, IcebergCatalogClient, IcebergTableId, LoadedIcebergTable, RestCatalogConfig,
9};
10
11// ── Unified Iceberg catalog backends (Phase J1-J4) ──────────────────────────
12// Each module is gated behind its own feature so a build that does not need a
13// given backend never pulls its dependency tree (sqlx, iceberg-catalog-rest…).
14#[cfg(feature = "glue-catalog")]
15pub mod glue_catalog;
16#[cfg(all(feature = "iceberg-datafusion", feature = "local-catalog"))]
17pub mod iceberg_catalog_bridge;
18#[cfg(feature = "iceberg-datafusion")]
19pub mod iceberg_table_provider;
20#[cfg(feature = "local-catalog")]
21pub mod local_catalog;
22
23pub mod object_store_io;
24#[cfg(feature = "postgres-catalog")]
25pub mod postgres_catalog;
26#[cfg(feature = "rest-catalog")]
27pub mod rest_catalog_wrapper;
28#[cfg(feature = "local-catalog")]
29pub mod unified;
30#[cfg(feature = "unity-catalog")]
31pub mod unity_catalog;
32
33#[cfg(feature = "glue-catalog")]
34pub use glue_catalog::GlueCatalog;
35#[cfg(all(feature = "iceberg-datafusion", feature = "local-catalog"))]
36pub use iceberg_catalog_bridge::IcebergCatalogBridge;
37#[cfg(feature = "local-catalog")]
38pub use local_catalog::LocalCatalog;
39#[cfg(feature = "postgres-catalog")]
40pub use postgres_catalog::PostgresCatalog;
41#[cfg(feature = "rest-catalog")]
42pub use rest_catalog_wrapper::KrishivRestCatalog;
43#[cfg(feature = "local-catalog")]
44pub use unified::KrishivCatalog;
45#[cfg(feature = "unity-catalog")]
46pub use unity_catalog::UnityCatalog;
47
48use std::collections::BTreeMap;
49use std::fmt;
50
51// ---------------------------------------------------------------------------
52// Error and Result
53// ---------------------------------------------------------------------------
54
55/// Errors produced by catalog operations.
56#[derive(Debug, thiserror::Error)]
57pub enum CatalogError {
58    /// A requested table was not found in the catalog.
59    #[error("table not found: '{name}'")]
60    TableNotFound { name: String },
61    /// Table already exists and `if_not_exists` was false.
62    #[error("table already exists: '{name}'")]
63    TableAlreadyExists { name: String },
64    /// A requested schema was not found.
65    #[error("schema not found: '{name}'")]
66    SchemaNotFound { name: String },
67    /// The provided schema is structurally invalid.
68    #[error("invalid schema: {message}")]
69    InvalidSchema { message: String },
70    /// Remote catalog configuration is malformed or unsafe.
71    #[error("invalid catalog configuration: {message}")]
72    InvalidConfiguration { message: String },
73    /// A remote catalog request could not be completed.
74    #[error("catalog transport error during {operation}: {message}")]
75    Transport { operation: String, message: String },
76    /// An HTTP request to a remote catalog service failed.
77    #[error("HTTP error {status}: {message}")]
78    Http { status: u16, message: String },
79    /// A successful remote response did not satisfy the catalog contract.
80    #[error("invalid catalog response during {operation}: {message}")]
81    InvalidResponse { operation: String, message: String },
82    /// A remote response exceeded the configured memory ceiling.
83    #[error("catalog response during {operation} exceeded {limit_bytes} bytes")]
84    ResponseTooLarge {
85        operation: String,
86        limit_bytes: usize,
87    },
88    /// The server explicitly does not advertise a required endpoint.
89    #[error("catalog server does not support {operation}")]
90    UnsupportedOperation { operation: String },
91    /// A filesystem / object-store I/O error.
92    #[error("I/O error: {0}")]
93    Io(String),
94    /// An error surfaced by the underlying Iceberg library.
95    #[error("Iceberg error: {0}")]
96    Iceberg(String),
97    /// An optimistic-concurrency commit lost a race with another writer.
98    #[error("concurrency conflict: {message}")]
99    ConcurrencyConflict { message: String },
100    /// A requested namespace was not found.
101    #[error("namespace not found: '{name}'")]
102    NamespaceNotFound { name: String },
103}
104
105/// Convenience result alias for catalog operations.
106pub type CatalogResult<T> = Result<T, CatalogError>;
107
108/// Alias used by the unified-catalog backends (J1-J4).
109///
110/// The file-system, Postgres, and REST catalog implementations all report
111/// failures through [`CatalogError`]; this alias gives those modules the
112/// `LakehouseError` name used throughout the lakehouse subsystem without
113/// introducing a second error type at the crate boundary.
114pub type LakehouseError = CatalogError;
115
116// ---------------------------------------------------------------------------
117// FieldType
118// ---------------------------------------------------------------------------
119
120/// Logical field types supported by the Krishiv catalog.
121#[derive(Debug, Clone, PartialEq, Eq)]
122pub enum FieldType {
123    Int8,
124    Int16,
125    Int32,
126    Int64,
127    UInt8,
128    UInt16,
129    UInt32,
130    UInt64,
131    Float32,
132    Float64,
133    Boolean,
134    Utf8,
135    Binary,
136    Timestamp,
137    Date32,
138    /// List of `item_type` elements.
139    List(Box<FieldType>),
140    /// Struct with named fields.
141    Struct(Vec<CatalogField>),
142}
143
144impl fmt::Display for FieldType {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        let s = match self {
147            FieldType::Int8 => "Int8",
148            FieldType::Int16 => "Int16",
149            FieldType::Int32 => "Int32",
150            FieldType::Int64 => "Int64",
151            FieldType::UInt8 => "UInt8",
152            FieldType::UInt16 => "UInt16",
153            FieldType::UInt32 => "UInt32",
154            FieldType::UInt64 => "UInt64",
155            FieldType::Float32 => "Float32",
156            FieldType::Float64 => "Float64",
157            FieldType::Boolean => "Boolean",
158            FieldType::Utf8 => "Utf8",
159            FieldType::Binary => "Binary",
160            FieldType::Timestamp => "Timestamp",
161            FieldType::Date32 => "Date32",
162            FieldType::List(inner) => return write!(f, "List<{inner}>"),
163            FieldType::Struct(fields) => {
164                return write!(f, "Struct({} fields)", fields.len());
165            }
166        };
167        f.write_str(s)
168    }
169}
170
171impl FieldType {
172    /// Convert this field type to the equivalent Arrow [`DataType`].
173    ///
174    /// [`DataType`]: arrow::datatypes::DataType
175    pub fn to_arrow(&self) -> arrow::datatypes::DataType {
176        use arrow::datatypes::DataType;
177        use arrow::datatypes::TimeUnit;
178        match self {
179            FieldType::Int8 => DataType::Int8,
180            FieldType::Int16 => DataType::Int16,
181            FieldType::Int32 => DataType::Int32,
182            FieldType::Int64 => DataType::Int64,
183            FieldType::UInt8 => DataType::UInt8,
184            FieldType::UInt16 => DataType::UInt16,
185            FieldType::UInt32 => DataType::UInt32,
186            FieldType::UInt64 => DataType::UInt64,
187            FieldType::Float32 => DataType::Float32,
188            FieldType::Float64 => DataType::Float64,
189            FieldType::Boolean => DataType::Boolean,
190            FieldType::Utf8 => DataType::Utf8,
191            FieldType::Binary => DataType::Binary,
192            FieldType::Timestamp => DataType::Timestamp(TimeUnit::Microsecond, None),
193            FieldType::Date32 => DataType::Date32,
194            FieldType::List(item) => DataType::List(std::sync::Arc::new(
195                arrow::datatypes::Field::new("item", item.to_arrow(), true),
196            )),
197            FieldType::Struct(fields) => {
198                let arrow_fields: arrow::datatypes::Fields = fields
199                    .iter()
200                    .map(|f| {
201                        std::sync::Arc::new(arrow::datatypes::Field::new(
202                            f.name(),
203                            f.field_type().to_arrow(),
204                            f.nullable(),
205                        ))
206                    })
207                    .collect();
208                DataType::Struct(arrow_fields)
209            }
210        }
211    }
212}
213
214// ---------------------------------------------------------------------------
215// CatalogField
216// ---------------------------------------------------------------------------
217
218/// A single field in a catalog table schema.
219#[derive(Debug, Clone, PartialEq, Eq)]
220pub struct CatalogField {
221    name: String,
222    field_type: FieldType,
223    nullable: bool,
224}
225
226impl CatalogField {
227    /// Create a new catalog field.
228    pub fn new(name: impl Into<String>, field_type: FieldType, nullable: bool) -> Self {
229        Self {
230            name: name.into(),
231            field_type,
232            nullable,
233        }
234    }
235
236    /// The field name.
237    pub fn name(&self) -> &str {
238        &self.name
239    }
240
241    /// The logical field type.
242    pub fn field_type(&self) -> &FieldType {
243        &self.field_type
244    }
245
246    /// Whether the field is nullable.
247    pub fn nullable(&self) -> bool {
248        self.nullable
249    }
250
251    /// Convert this field to an Arrow [`Field`].
252    ///
253    /// [`Field`]: arrow::datatypes::Field
254    pub fn to_arrow_field(&self) -> arrow::datatypes::Field {
255        arrow::datatypes::Field::new(
256            self.name.as_str(),
257            self.field_type.to_arrow(),
258            self.nullable,
259        )
260    }
261}
262
263// ---------------------------------------------------------------------------
264// TableSchema
265// ---------------------------------------------------------------------------
266
267/// The schema of a catalog table: an ordered list of fields.
268#[derive(Debug, Clone, PartialEq, Eq)]
269pub struct TableSchema {
270    fields: Vec<CatalogField>,
271}
272
273impl TableSchema {
274    /// Create a new schema from a list of fields.
275    pub fn new(fields: Vec<CatalogField>) -> Self {
276        Self { fields }
277    }
278
279    /// Create an empty schema with no fields.
280    pub fn empty() -> Self {
281        Self { fields: Vec::new() }
282    }
283
284    /// Convert to an Arrow [`Schema`].
285    ///
286    /// [`Schema`]: arrow::datatypes::Schema
287    pub fn to_arrow_schema(&self) -> arrow::datatypes::Schema {
288        let arrow_fields: Vec<arrow::datatypes::Field> = self
289            .fields
290            .iter()
291            .map(CatalogField::to_arrow_field)
292            .collect();
293        arrow::datatypes::Schema::new(arrow_fields)
294    }
295
296    /// Return the number of fields in this schema.
297    pub fn field_count(&self) -> usize {
298        self.fields.len()
299    }
300
301    /// Look up a field by name.
302    pub fn get_field(&self, name: &str) -> Option<&CatalogField> {
303        self.fields.iter().find(|f| f.name() == name)
304    }
305}
306
307// ---------------------------------------------------------------------------
308// ColumnStatistics
309// ---------------------------------------------------------------------------
310
311/// Optional statistics for a table column (or the whole table).
312#[derive(Debug, Clone, PartialEq, Eq, Default)]
313pub struct ColumnStatistics {
314    /// Total number of rows in the table, if known.
315    pub row_count: Option<u64>,
316    /// Number of null values in the column, if known.
317    pub null_count: Option<u64>,
318    /// String representation of the minimum value, if known.
319    pub min_value: Option<String>,
320    /// String representation of the maximum value, if known.
321    pub max_value: Option<String>,
322    /// Number of distinct values (NDV) for the column, if known.
323    ///
324    /// Drives join-order cost estimation and broadcast sizing: a column
325    /// with NDV ≈ 1 is a hot join key, while a column with NDV ≈ row
326    /// count is a near-unique key (build-side fan-out is high).
327    pub distinct_count: Option<u64>,
328    /// Wall-clock Unix timestamp (seconds) when these stats were
329    /// collected by `ANALYZE TABLE`. The CBO refuses to use stats
330    /// older than `stats_max_age_secs` to avoid optimising against
331    /// stale cardinality information.
332    pub collected_at_secs: Option<u64>,
333}
334
335impl ColumnStatistics {
336    /// Create a new `ColumnStatistics` with all fields set to `None`.
337    pub fn new() -> Self {
338        Self::default()
339    }
340
341    /// Set the row count.
342    #[must_use]
343    pub fn with_row_count(mut self, count: u64) -> Self {
344        self.row_count = Some(count);
345        self
346    }
347
348    /// Set the null count.
349    #[must_use]
350    pub fn with_null_count(mut self, count: u64) -> Self {
351        self.null_count = Some(count);
352        self
353    }
354
355    /// Set the minimum value.
356    #[must_use]
357    pub fn with_min(mut self, min: impl Into<String>) -> Self {
358        self.min_value = Some(min.into());
359        self
360    }
361
362    /// Set the maximum value.
363    #[must_use]
364    pub fn with_max(mut self, max: impl Into<String>) -> Self {
365        self.max_value = Some(max.into());
366        self
367    }
368
369    /// Set the NDV (number of distinct values) for the column.
370    #[must_use]
371    pub fn with_distinct_count(mut self, ndv: u64) -> Self {
372        self.distinct_count = Some(ndv);
373        self
374    }
375
376    /// Stamp the wall-clock timestamp (seconds) at which these stats
377    /// were collected. Stamped by [`analyze_table`][crate::analyze::analyze_table]
378    /// so the CBO can detect stale stats.
379    #[must_use]
380    pub fn with_collected_at_secs(mut self, secs: u64) -> Self {
381        self.collected_at_secs = Some(secs);
382        self
383    }
384
385    /// Selectivity of a single-point equality filter on this column.
386    ///
387    /// `1.0 / distinct_count` when both are known; otherwise `None`.
388    /// `0.0` when `distinct_count == 0` (the column is all-null or empty).
389    pub fn equality_selectivity(&self) -> Option<f64> {
390        let ndv = self.distinct_count?;
391        if ndv == 0 {
392            Some(0.0)
393        } else {
394            Some(1.0 / ndv as f64)
395        }
396    }
397
398    /// Whether the stats are "fresh" given a current timestamp and a
399    /// maximum age in seconds.
400    ///
401    /// Stats without a collection timestamp are considered fresh by
402    /// default (back-compat for manually-curated stats). Stats older
403    /// than `max_age_secs` are stale.
404    pub fn is_fresh(&self, now_secs: u64, max_age_secs: u64) -> bool {
405        match self.collected_at_secs {
406            None => true,
407            Some(ts) => now_secs.saturating_sub(ts) <= max_age_secs,
408        }
409    }
410}
411
412// ---------------------------------------------------------------------------
413// TableMetadata
414// ---------------------------------------------------------------------------
415
416/// Full metadata for a table: name, schema, and optional statistics.
417#[derive(Debug, Clone)]
418pub struct TableMetadata {
419    name: String,
420    schema: TableSchema,
421    stats: Option<ColumnStatistics>,
422}
423
424impl TableMetadata {
425    /// Create new table metadata with no statistics.
426    pub fn new(name: impl Into<String>, schema: TableSchema) -> Self {
427        Self {
428            name: name.into(),
429            schema,
430            stats: None,
431        }
432    }
433
434    /// Attach column statistics and return the updated metadata.
435    #[must_use]
436    pub fn with_stats(mut self, stats: ColumnStatistics) -> Self {
437        self.stats = Some(stats);
438        self
439    }
440
441    /// The table name.
442    pub fn name(&self) -> &str {
443        &self.name
444    }
445
446    /// The table schema.
447    pub fn schema(&self) -> &TableSchema {
448        &self.schema
449    }
450
451    /// Optional column statistics.
452    pub fn statistics(&self) -> Option<&ColumnStatistics> {
453        self.stats.as_ref()
454    }
455}
456
457// ---------------------------------------------------------------------------
458// TableProvider trait
459// ---------------------------------------------------------------------------
460
461/// A resolved reference to a single table's metadata.
462pub trait TableProvider {
463    /// The table name.
464    fn name(&self) -> &str;
465
466    /// The table schema.
467    fn schema(&self) -> &TableSchema;
468
469    /// Optional column statistics.
470    fn statistics(&self) -> Option<&ColumnStatistics>;
471}
472
473// ---------------------------------------------------------------------------
474// CatalogProvider trait
475// ---------------------------------------------------------------------------
476
477/// A registry of tables that can be listed, looked up, and registered.
478pub trait CatalogProvider {
479    /// Return the names of all tables in the catalog.
480    fn list_tables(&self) -> Vec<String>;
481
482    /// Look up a table by name.
483    fn get_table(&self, name: &str) -> CatalogResult<&dyn TableProvider>;
484
485    /// Register a table in the catalog.
486    ///
487    /// Returns an error if the schema is structurally invalid or if
488    /// implementation-specific constraints are violated.
489    fn register_table(&mut self, metadata: TableMetadata) -> CatalogResult<()>;
490}
491
492// ---------------------------------------------------------------------------
493// InMemoryCatalog
494// ---------------------------------------------------------------------------
495
496/// A `TableProvider` wrapper over `TableMetadata`.
497struct TableMetadataProvider {
498    metadata: TableMetadata,
499}
500
501impl TableProvider for TableMetadataProvider {
502    fn name(&self) -> &str {
503        self.metadata.name()
504    }
505
506    fn schema(&self) -> &TableSchema {
507        self.metadata.schema()
508    }
509
510    fn statistics(&self) -> Option<&ColumnStatistics> {
511        self.metadata.statistics()
512    }
513}
514
515/// An in-memory catalog backed by a sorted map.
516pub struct InMemoryCatalog {
517    tables: BTreeMap<String, TableMetadataProvider>,
518    /// Optional in-memory row data keyed by table name.
519    table_data: BTreeMap<String, std::sync::Arc<Vec<arrow::record_batch::RecordBatch>>>,
520}
521
522impl InMemoryCatalog {
523    /// Create a new, empty in-memory catalog.
524    pub fn new() -> Self {
525        Self {
526            tables: BTreeMap::new(),
527            table_data: BTreeMap::new(),
528        }
529    }
530
531    /// Register a table and attach in-memory Arrow batches for SQL scans (P0-9).
532    pub fn register_table_with_batches(
533        &mut self,
534        metadata: TableMetadata,
535        batches: Vec<arrow::record_batch::RecordBatch>,
536    ) -> CatalogResult<()> {
537        let name = metadata.name().to_owned();
538        self.register_table(metadata)?;
539        if !batches.is_empty() {
540            self.table_data.insert(name, std::sync::Arc::new(batches));
541        }
542        Ok(())
543    }
544
545    /// Return stored batches for a registered table, if any.
546    pub fn table_batches(
547        &self,
548        name: &str,
549    ) -> Option<std::sync::Arc<Vec<arrow::record_batch::RecordBatch>>> {
550        self.table_data.get(name).cloned()
551    }
552}
553
554impl Default for InMemoryCatalog {
555    fn default() -> Self {
556        Self::new()
557    }
558}
559
560impl CatalogProvider for InMemoryCatalog {
561    fn list_tables(&self) -> Vec<String> {
562        self.tables.keys().cloned().collect()
563    }
564
565    fn get_table(&self, name: &str) -> CatalogResult<&dyn TableProvider> {
566        self.tables
567            .get(name)
568            .map(|p| p as &dyn TableProvider)
569            .ok_or_else(|| CatalogError::TableNotFound {
570                name: name.to_string(),
571            })
572    }
573
574    fn register_table(&mut self, metadata: TableMetadata) -> CatalogResult<()> {
575        let name = metadata.name().to_string();
576        if self.tables.contains_key(&name) {
577            return Err(CatalogError::TableAlreadyExists { name });
578        }
579        self.tables.insert(name, TableMetadataProvider { metadata });
580        Ok(())
581    }
582}
583
584// ---------------------------------------------------------------------------
585// SchemaRegistry
586// ---------------------------------------------------------------------------
587
588/// A registry that maps logical schema names to [`TableSchema`] definitions.
589///
590/// Used by connectors (e.g., Kafka Avro/Protobuf topics) to resolve the
591/// Arrow schema for a data stream at runtime without hard-coding field lists.
592pub trait SchemaRegistry: Send + Sync {
593    /// Look up a schema by name.
594    fn get_schema(&self, name: &str) -> CatalogResult<TableSchema>;
595    /// Register a schema under a name, replacing any existing entry.
596    fn register_schema(&mut self, name: impl Into<String>, schema: TableSchema);
597    /// Return all registered schema names.
598    fn schema_names(&self) -> Vec<String>;
599}
600
601/// An in-memory [`SchemaRegistry`] backed by a sorted map.
602#[derive(Debug, Default)]
603pub struct InMemorySchemaRegistry {
604    schemas: BTreeMap<String, TableSchema>,
605}
606
607impl InMemorySchemaRegistry {
608    pub fn new() -> Self {
609        Self::default()
610    }
611}
612
613impl SchemaRegistry for InMemorySchemaRegistry {
614    fn get_schema(&self, name: &str) -> CatalogResult<TableSchema> {
615        self.schemas
616            .get(name)
617            .cloned()
618            .ok_or_else(|| CatalogError::SchemaNotFound {
619                name: name.to_string(),
620            })
621    }
622
623    fn register_schema(&mut self, name: impl Into<String>, schema: TableSchema) {
624        self.schemas.insert(name.into(), schema);
625    }
626
627    fn schema_names(&self) -> Vec<String> {
628        self.schemas.keys().cloned().collect()
629    }
630}
631
632// ---------------------------------------------------------------------------
633// DataFusion catalog bridge
634// ---------------------------------------------------------------------------
635
636/// DataFusion integration: wraps [`InMemoryCatalog`] as DataFusion catalog
637/// and schema providers so that Krishiv catalog tables can be used directly
638/// inside a DataFusion [`SessionContext`].
639///
640/// [`SessionContext`]: datafusion::prelude::SessionContext
641pub mod datafusion_bridge {
642
643    use std::fmt;
644    use std::sync::{Arc, RwLock};
645
646    use datafusion::catalog::{CatalogProvider, SchemaProvider};
647    use datafusion::datasource::MemTable;
648    use datafusion::error::Result as DfResult;
649
650    /// Bridges a Krishiv [`InMemoryCatalog`] into a DataFusion
651    /// [`CatalogProvider`].
652    ///
653    /// The bridge exposes a single schema named `"public"` that mirrors
654    /// the tables registered in the underlying [`InMemoryCatalog`].
655    ///
656    /// [`InMemoryCatalog`]: super::InMemoryCatalog
657    pub struct DataFusionCatalogBridge {
658        catalog: Arc<RwLock<super::InMemoryCatalog>>,
659        schema_name: String,
660        /// MemTable cache shared across the inner `DataFusionSchemaBridge`
661        /// instances DataFusion requests. Avoids re-cloning the entire
662        /// `Vec<RecordBatch>` payload (which can be hundreds of MB) on
663        /// every DataFusion `table()` call. Keyed by table name; cleared
664        /// on `invalidate(name)`.
665        schema_cache: std::sync::Arc<dashmap::DashMap<String, Arc<MemTable>>>,
666    }
667
668    impl DataFusionCatalogBridge {
669        /// Create a bridge from an [`InMemoryCatalog`] shared reference.
670        ///
671        /// [`InMemoryCatalog`]: super::InMemoryCatalog
672        pub fn new(catalog: Arc<RwLock<super::InMemoryCatalog>>) -> Self {
673            Self {
674                catalog,
675                schema_name: "public".to_string(),
676                schema_cache: std::sync::Arc::new(dashmap::DashMap::new()),
677            }
678        }
679
680        /// Invalidate the MemTable cache for `name`, forcing the next
681        /// `table()` call to rebuild the cached `MemTable` from the
682        /// catalog's batch store. Call this after `register_table_with_batches`
683        /// mutates the underlying catalog so the bridge does not serve
684        /// a stale `MemTable` referencing the previous batch payload.
685        ///
686        /// The `DashMap` cache is per-bridge and survives across DataFusion
687        /// `table()` calls; without this invalidation hook a second
688        /// `register_table_with_batches` for the same name would not be
689        /// visible to the DataFusion query plan.
690        pub fn invalidate(&self, name: &str) {
691            self.schema_cache.remove(name);
692        }
693    }
694
695    impl fmt::Debug for DataFusionCatalogBridge {
696        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
697            f.debug_struct("DataFusionCatalogBridge")
698                .field("schema_name", &self.schema_name)
699                .finish()
700        }
701    }
702
703    impl CatalogProvider for DataFusionCatalogBridge {
704        fn schema_names(&self) -> Vec<String> {
705            vec![self.schema_name.clone()]
706        }
707
708        fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
709            if name == self.schema_name {
710                Some(Arc::new(DataFusionSchemaBridge {
711                    catalog: self.catalog.clone(),
712                    cache: self.schema_cache.clone(),
713                }))
714            } else {
715                None
716            }
717        }
718    }
719
720    // -----------------------------------------------------------------------
721
722    struct DataFusionSchemaBridge {
723        catalog: Arc<RwLock<super::InMemoryCatalog>>,
724        cache: std::sync::Arc<dashmap::DashMap<String, Arc<MemTable>>>,
725    }
726
727    impl fmt::Debug for DataFusionSchemaBridge {
728        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
729            f.debug_struct("DataFusionSchemaBridge").finish()
730        }
731    }
732
733    #[async_trait::async_trait]
734    impl SchemaProvider for DataFusionSchemaBridge {
735        fn table_names(&self) -> Vec<String> {
736            let catalog = self.catalog.read().unwrap_or_else(|p| p.into_inner());
737            use super::CatalogProvider as KrishivCatalogProvider;
738            catalog.list_tables()
739        }
740
741        async fn table(
742            &self,
743            name: &str,
744        ) -> DfResult<Option<Arc<dyn datafusion::datasource::TableProvider>>> {
745            // MemTable cache: if we built one for this table already, return
746            // the cached Arc. The cache is invalidated explicitly via
747            // `invalidate(name)` when the underlying catalog is mutated.
748            if let Some(cached) = self.cache.get(name) {
749                return Ok(Some(
750                    cached.clone() as Arc<dyn datafusion::datasource::TableProvider>
751                ));
752            }
753            let catalog = self.catalog.read().unwrap_or_else(|p| p.into_inner());
754            use super::CatalogProvider as KrishivCatalogProvider;
755            match catalog.get_table(name) {
756                Ok(table_provider) => {
757                    let arrow_schema = Arc::new(table_provider.schema().to_arrow_schema());
758                    let batches = catalog.table_batches(name);
759                    let partitions = batches.map(|b| (*b).clone()).unwrap_or_default();
760                    let mem = MemTable::try_new(arrow_schema, vec![partitions])?;
761                    let mem_arc = Arc::new(mem);
762                    self.cache.insert(name.to_string(), mem_arc.clone());
763                    Ok(Some(
764                        mem_arc as Arc<dyn datafusion::datasource::TableProvider>,
765                    ))
766                }
767                Err(super::CatalogError::TableNotFound { .. }) => Ok(None),
768                Err(error) => Err(datafusion::error::DataFusionError::External(Box::new(
769                    error,
770                ))),
771            }
772        }
773
774        fn table_exist(&self, name: &str) -> bool {
775            let catalog = self.catalog.read().unwrap_or_else(|p| p.into_inner());
776            use super::CatalogProvider as KrishivCatalogProvider;
777            catalog.get_table(name).is_ok()
778        }
779    }
780}
781
782// ---------------------------------------------------------------------------
783// Unit tests
784// ---------------------------------------------------------------------------
785
786#[cfg(test)]
787mod tests {
788    use super::*;
789
790    fn make_schema() -> TableSchema {
791        TableSchema::new(vec![
792            CatalogField::new("id", FieldType::Int64, false),
793            CatalogField::new("name", FieldType::Utf8, true),
794        ])
795    }
796
797    #[test]
798    fn in_memory_catalog_registers_and_retrieves_table() {
799        let mut catalog = InMemoryCatalog::new();
800        let meta = TableMetadata::new("users", make_schema());
801        catalog.register_table(meta).unwrap();
802
803        let table = catalog.get_table("users").unwrap();
804        assert_eq!(table.name(), "users");
805        assert_eq!(table.schema().field_count(), 2);
806    }
807
808    #[test]
809    fn in_memory_catalog_lists_tables() {
810        let mut catalog = InMemoryCatalog::new();
811        catalog
812            .register_table(TableMetadata::new("alpha", make_schema()))
813            .unwrap();
814        catalog
815            .register_table(TableMetadata::new("beta", make_schema()))
816            .unwrap();
817
818        let mut tables = catalog.list_tables();
819        tables.sort();
820        assert_eq!(tables, vec!["alpha", "beta"]);
821    }
822
823    #[test]
824    fn in_memory_catalog_returns_error_for_unknown_table() {
825        let catalog = InMemoryCatalog::new();
826        let err = catalog.get_table("nonexistent").err().unwrap();
827        match err {
828            CatalogError::TableNotFound { name } => {
829                assert_eq!(name, "nonexistent");
830            }
831            other => panic!("unexpected error: {other}"),
832        }
833    }
834
835    #[test]
836    fn table_schema_converts_to_arrow_schema() {
837        let schema = make_schema();
838        let arrow_schema = schema.to_arrow_schema();
839
840        assert_eq!(arrow_schema.fields().len(), 2);
841        let id_field = arrow_schema.field_with_name("id").unwrap();
842        assert_eq!(id_field.data_type(), &arrow::datatypes::DataType::Int64);
843        assert!(!id_field.is_nullable());
844
845        let name_field = arrow_schema.field_with_name("name").unwrap();
846        assert_eq!(name_field.data_type(), &arrow::datatypes::DataType::Utf8);
847        assert!(name_field.is_nullable());
848    }
849
850    // -----------------------------------------------------------------------
851    // SchemaRegistry tests
852    // -----------------------------------------------------------------------
853
854    #[test]
855    fn schema_registry_registers_and_retrieves() {
856        let mut registry = InMemorySchemaRegistry::new();
857        registry.register_schema("events", make_schema());
858        let schema = registry.get_schema("events").unwrap();
859        assert_eq!(schema.field_count(), 2);
860    }
861
862    #[test]
863    fn schema_registry_returns_error_for_missing() {
864        let registry = InMemorySchemaRegistry::new();
865        let err = registry.get_schema("nonexistent").unwrap_err();
866        match err {
867            CatalogError::SchemaNotFound { name } => {
868                assert_eq!(name, "nonexistent");
869            }
870            other => panic!("unexpected error: {other}"),
871        }
872    }
873
874    #[test]
875    fn schema_registry_lists_names() {
876        let mut registry = InMemorySchemaRegistry::new();
877        registry.register_schema("orders", make_schema());
878        registry.register_schema("users", make_schema());
879        let mut names = registry.schema_names();
880        names.sort();
881        assert_eq!(names, vec!["orders", "users"]);
882    }
883
884    // -----------------------------------------------------------------------
885    // DataFusion bridge tests
886    // -----------------------------------------------------------------------
887
888    #[test]
889    fn datafusion_bridge_schema_names_returns_public() {
890        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
891
892        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
893        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
894        let names = bridge.schema_names();
895        assert_eq!(names, vec!["public"]);
896    }
897
898    #[test]
899    fn datafusion_bridge_table_exist() {
900        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
901        {
902            let mut cat = catalog.write().unwrap();
903            cat.register_table(TableMetadata::new("orders", make_schema()))
904                .unwrap();
905        }
906        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
907        let schema_provider = {
908            use datafusion::catalog::CatalogProvider as DfCatalogProvider;
909            bridge.schema("public").unwrap()
910        };
911        assert!(schema_provider.table_exist("orders"));
912        assert!(!schema_provider.table_exist("nonexistent"));
913    }
914
915    #[tokio::test]
916    async fn datafusion_bridge_memtable_cache_reuses_arc() {
917        use std::sync::Arc;
918
919        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
920        {
921            let mut cat = catalog.write().unwrap();
922            cat.register_table(TableMetadata::new("orders", make_schema()))
923                .unwrap();
924        }
925        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
926        let schema_provider = {
927            use datafusion::catalog::CatalogProvider as DfCatalogProvider;
928            bridge.schema("public").unwrap()
929        };
930        let first = schema_provider.table("orders").await.unwrap().unwrap();
931        let second = schema_provider.table("orders").await.unwrap().unwrap();
932        // Cached: identical Arc pointer, no re-clone of batch payload.
933        let cached = Arc::ptr_eq(&first, &second);
934        assert!(cached, "expected cached MemTable Arc, got fresh allocation");
935    }
936
937    #[tokio::test]
938    async fn datafusion_bridge_invalidate_forces_rebuild() {
939        use std::sync::Arc;
940
941        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
942        {
943            let mut cat = catalog.write().unwrap();
944            cat.register_table(TableMetadata::new("orders", make_schema()))
945                .unwrap();
946        }
947        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
948        let schema_provider = {
949            use datafusion::catalog::CatalogProvider as DfCatalogProvider;
950            bridge.schema("public").unwrap()
951        };
952        let first = schema_provider.table("orders").await.unwrap().unwrap();
953        bridge.invalidate("orders");
954        let second = schema_provider.table("orders").await.unwrap().unwrap();
955        // After invalidation, a new MemTable Arc is constructed.
956        assert!(!Arc::ptr_eq(&first, &second));
957        // A second call without further invalidation must hit the cache again.
958        let third = schema_provider.table("orders").await.unwrap().unwrap();
959        assert!(Arc::ptr_eq(&second, &third));
960    }
961
962    #[tokio::test]
963    async fn catalog_scan_returns_registered_row_count() {
964        use std::sync::Arc;
965
966        use arrow::array::Int64Array;
967        use arrow::datatypes::{DataType, Field, Schema};
968        use arrow::record_batch::RecordBatch;
969        use datafusion::prelude::SessionContext;
970
971        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
972        let schema = TableSchema::new(vec![CatalogField::new("id", FieldType::Int64, false)]);
973        let arrow_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
974        let values: Vec<Option<i64>> = (0..10).map(Some).collect();
975        let batch =
976            RecordBatch::try_new(arrow_schema, vec![Arc::new(Int64Array::from(values))]).unwrap();
977        catalog
978            .write()
979            .unwrap()
980            .register_table_with_batches(TableMetadata::new("t", schema), vec![batch])
981            .unwrap();
982
983        let ctx = SessionContext::new();
984        ctx.register_catalog(
985            "krishiv",
986            Arc::new(super::datafusion_bridge::DataFusionCatalogBridge::new(
987                catalog,
988            )),
989        );
990        let df = ctx.sql("SELECT * FROM krishiv.public.t").await.unwrap();
991        let batches = df.collect().await.unwrap();
992        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
993        assert_eq!(rows, 10);
994    }
995
996    #[test]
997    fn datafusion_bridge_unknown_schema_returns_none() {
998        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
999
1000        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
1001        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
1002        let result = bridge.schema("nonexistent");
1003        assert!(result.is_none());
1004    }
1005
1006    // -----------------------------------------------------------------------
1007    // CatalogError Display tests
1008    // -----------------------------------------------------------------------
1009
1010    #[test]
1011    fn catalog_error_display_table_not_found() {
1012        let err = CatalogError::TableNotFound {
1013            name: "orders".to_string(),
1014        };
1015        assert_eq!(err.to_string(), "table not found: 'orders'");
1016    }
1017
1018    #[test]
1019    fn catalog_error_display_table_already_exists() {
1020        let err = CatalogError::TableAlreadyExists {
1021            name: "users".to_string(),
1022        };
1023        assert_eq!(err.to_string(), "table already exists: 'users'");
1024    }
1025
1026    #[test]
1027    fn catalog_error_display_schema_not_found() {
1028        let err = CatalogError::SchemaNotFound {
1029            name: "events".to_string(),
1030        };
1031        assert_eq!(err.to_string(), "schema not found: 'events'");
1032    }
1033
1034    #[test]
1035    fn catalog_error_display_invalid_schema() {
1036        let err = CatalogError::InvalidSchema {
1037            message: "missing required field 'id'".to_string(),
1038        };
1039        assert_eq!(
1040            err.to_string(),
1041            "invalid schema: missing required field 'id'"
1042        );
1043    }
1044
1045    #[test]
1046    fn catalog_error_display_http() {
1047        let err = CatalogError::Http {
1048            status: 404,
1049            message: "not found".to_string(),
1050        };
1051        assert_eq!(err.to_string(), "HTTP error 404: not found");
1052    }
1053
1054    #[test]
1055    fn catalog_error_is_std_error() {
1056        let err = CatalogError::TableNotFound {
1057            name: "t".to_string(),
1058        };
1059        let e: &dyn std::error::Error = &err;
1060        assert!(e.source().is_none());
1061    }
1062
1063    // -----------------------------------------------------------------------
1064    // FieldType to_arrow tests
1065    // -----------------------------------------------------------------------
1066
1067    #[test]
1068    fn field_type_to_arrow_int8() {
1069        assert_eq!(FieldType::Int8.to_arrow(), arrow::datatypes::DataType::Int8);
1070    }
1071
1072    #[test]
1073    fn field_type_to_arrow_int16() {
1074        assert_eq!(
1075            FieldType::Int16.to_arrow(),
1076            arrow::datatypes::DataType::Int16
1077        );
1078    }
1079
1080    #[test]
1081    fn field_type_to_arrow_int32() {
1082        assert_eq!(
1083            FieldType::Int32.to_arrow(),
1084            arrow::datatypes::DataType::Int32
1085        );
1086    }
1087
1088    #[test]
1089    fn field_type_to_arrow_int64() {
1090        assert_eq!(
1091            FieldType::Int64.to_arrow(),
1092            arrow::datatypes::DataType::Int64
1093        );
1094    }
1095
1096    #[test]
1097    fn field_type_to_arrow_uint8() {
1098        assert_eq!(
1099            FieldType::UInt8.to_arrow(),
1100            arrow::datatypes::DataType::UInt8
1101        );
1102    }
1103
1104    #[test]
1105    fn field_type_to_arrow_uint16() {
1106        assert_eq!(
1107            FieldType::UInt16.to_arrow(),
1108            arrow::datatypes::DataType::UInt16
1109        );
1110    }
1111
1112    #[test]
1113    fn field_type_to_arrow_uint32() {
1114        assert_eq!(
1115            FieldType::UInt32.to_arrow(),
1116            arrow::datatypes::DataType::UInt32
1117        );
1118    }
1119
1120    #[test]
1121    fn field_type_to_arrow_uint64() {
1122        assert_eq!(
1123            FieldType::UInt64.to_arrow(),
1124            arrow::datatypes::DataType::UInt64
1125        );
1126    }
1127
1128    #[test]
1129    fn field_type_to_arrow_float32() {
1130        assert_eq!(
1131            FieldType::Float32.to_arrow(),
1132            arrow::datatypes::DataType::Float32
1133        );
1134    }
1135
1136    #[test]
1137    fn field_type_to_arrow_float64() {
1138        assert_eq!(
1139            FieldType::Float64.to_arrow(),
1140            arrow::datatypes::DataType::Float64
1141        );
1142    }
1143
1144    #[test]
1145    fn field_type_to_arrow_boolean() {
1146        assert_eq!(
1147            FieldType::Boolean.to_arrow(),
1148            arrow::datatypes::DataType::Boolean
1149        );
1150    }
1151
1152    #[test]
1153    fn field_type_to_arrow_utf8() {
1154        assert_eq!(FieldType::Utf8.to_arrow(), arrow::datatypes::DataType::Utf8);
1155    }
1156
1157    #[test]
1158    fn field_type_to_arrow_binary() {
1159        assert_eq!(
1160            FieldType::Binary.to_arrow(),
1161            arrow::datatypes::DataType::Binary
1162        );
1163    }
1164
1165    #[test]
1166    fn field_type_to_arrow_timestamp() {
1167        use arrow::datatypes::{DataType, TimeUnit};
1168        assert_eq!(
1169            FieldType::Timestamp.to_arrow(),
1170            DataType::Timestamp(TimeUnit::Microsecond, None)
1171        );
1172    }
1173
1174    #[test]
1175    fn field_type_to_arrow_date32() {
1176        assert_eq!(
1177            FieldType::Date32.to_arrow(),
1178            arrow::datatypes::DataType::Date32
1179        );
1180    }
1181
1182    #[test]
1183    fn field_type_to_arrow_list() {
1184        let list_type = FieldType::List(Box::new(FieldType::Utf8));
1185        match list_type.to_arrow() {
1186            arrow::datatypes::DataType::List(field) => {
1187                assert_eq!(field.name(), "item");
1188                assert_eq!(field.data_type(), &arrow::datatypes::DataType::Utf8);
1189                assert!(field.is_nullable());
1190            }
1191            other => panic!("expected List, got {other:?}"),
1192        }
1193    }
1194
1195    #[test]
1196    fn field_type_to_arrow_struct() {
1197        let struct_type = FieldType::Struct(vec![
1198            CatalogField::new("x", FieldType::Int32, false),
1199            CatalogField::new("y", FieldType::Utf8, true),
1200        ]);
1201        match struct_type.to_arrow() {
1202            arrow::datatypes::DataType::Struct(fields) => {
1203                assert_eq!(fields.len(), 2);
1204                assert_eq!(fields[0].name(), "x");
1205                assert_eq!(fields[0].data_type(), &arrow::datatypes::DataType::Int32);
1206                assert!(!fields[0].is_nullable());
1207                assert_eq!(fields[1].name(), "y");
1208                assert_eq!(fields[1].data_type(), &arrow::datatypes::DataType::Utf8);
1209                assert!(fields[1].is_nullable());
1210            }
1211            other => panic!("expected Struct, got {other:?}"),
1212        }
1213    }
1214
1215    // -----------------------------------------------------------------------
1216    // FieldType Display tests
1217    // -----------------------------------------------------------------------
1218
1219    #[test]
1220    fn field_type_display_simple() {
1221        assert_eq!(FieldType::Int8.to_string(), "Int8");
1222        assert_eq!(FieldType::Int16.to_string(), "Int16");
1223        assert_eq!(FieldType::Int32.to_string(), "Int32");
1224        assert_eq!(FieldType::Int64.to_string(), "Int64");
1225        assert_eq!(FieldType::UInt8.to_string(), "UInt8");
1226        assert_eq!(FieldType::UInt16.to_string(), "UInt16");
1227        assert_eq!(FieldType::UInt32.to_string(), "UInt32");
1228        assert_eq!(FieldType::UInt64.to_string(), "UInt64");
1229        assert_eq!(FieldType::Float32.to_string(), "Float32");
1230        assert_eq!(FieldType::Float64.to_string(), "Float64");
1231        assert_eq!(FieldType::Boolean.to_string(), "Boolean");
1232        assert_eq!(FieldType::Utf8.to_string(), "Utf8");
1233        assert_eq!(FieldType::Binary.to_string(), "Binary");
1234        assert_eq!(FieldType::Timestamp.to_string(), "Timestamp");
1235        assert_eq!(FieldType::Date32.to_string(), "Date32");
1236    }
1237
1238    #[test]
1239    fn field_type_display_list() {
1240        let list = FieldType::List(Box::new(FieldType::Int32));
1241        assert_eq!(list.to_string(), "List<Int32>");
1242    }
1243
1244    #[test]
1245    fn field_type_display_struct() {
1246        let s = FieldType::Struct(vec![
1247            CatalogField::new("a", FieldType::Boolean, true),
1248            CatalogField::new("b", FieldType::Utf8, false),
1249        ]);
1250        assert_eq!(s.to_string(), "Struct(2 fields)");
1251    }
1252
1253    // -----------------------------------------------------------------------
1254    // CatalogField tests
1255    // -----------------------------------------------------------------------
1256
1257    #[test]
1258    fn catalog_field_accessors() {
1259        let f = CatalogField::new("col", FieldType::Float64, true);
1260        assert_eq!(f.name(), "col");
1261        assert_eq!(f.field_type(), &FieldType::Float64);
1262        assert!(f.nullable());
1263    }
1264
1265    #[test]
1266    fn catalog_field_to_arrow_field() {
1267        let f = CatalogField::new("ts", FieldType::Timestamp, false);
1268        let arrow_f = f.to_arrow_field();
1269        assert_eq!(arrow_f.name(), "ts");
1270        use arrow::datatypes::{DataType, TimeUnit};
1271        assert_eq!(
1272            arrow_f.data_type(),
1273            &DataType::Timestamp(TimeUnit::Microsecond, None)
1274        );
1275        assert!(!arrow_f.is_nullable());
1276    }
1277
1278    // -----------------------------------------------------------------------
1279    // ColumnStatistics tests
1280    // -----------------------------------------------------------------------
1281
1282    #[test]
1283    fn column_statistics_new_defaults() {
1284        let stats = ColumnStatistics::new();
1285        assert!(stats.row_count.is_none());
1286        assert!(stats.null_count.is_none());
1287        assert!(stats.min_value.is_none());
1288        assert!(stats.max_value.is_none());
1289    }
1290
1291    #[test]
1292    fn column_statistics_default_trait() {
1293        let stats = ColumnStatistics::default();
1294        assert_eq!(stats, ColumnStatistics::new());
1295    }
1296
1297    #[test]
1298    fn column_statistics_builder_all_fields() {
1299        let stats = ColumnStatistics::new()
1300            .with_row_count(1_000_000)
1301            .with_null_count(42)
1302            .with_min("abc")
1303            .with_max("xyz");
1304
1305        assert_eq!(stats.row_count, Some(1_000_000));
1306        assert_eq!(stats.null_count, Some(42));
1307        assert_eq!(stats.min_value.as_deref(), Some("abc"));
1308        assert_eq!(stats.max_value.as_deref(), Some("xyz"));
1309    }
1310
1311    #[test]
1312    fn column_statistics_builder_partial() {
1313        let stats = ColumnStatistics::new().with_row_count(500);
1314        assert_eq!(stats.row_count, Some(500));
1315        assert!(stats.null_count.is_none());
1316        assert!(stats.min_value.is_none());
1317        assert!(stats.max_value.is_none());
1318    }
1319
1320    #[test]
1321    fn column_statistics_into_string() {
1322        let stats = ColumnStatistics::new()
1323            .with_row_count(100)
1324            .with_null_count(5)
1325            .with_min("1")
1326            .with_max("99");
1327        let dbg = format!("{stats:?}");
1328        assert!(dbg.contains("row_count: Some(100)"));
1329        assert!(dbg.contains("null_count: Some(5)"));
1330        assert!(dbg.contains("min_value: Some(\"1\")"));
1331        assert!(dbg.contains("max_value: Some(\"99\")"));
1332    }
1333
1334    #[test]
1335    fn column_statistics_builder_overwrites() {
1336        let stats = ColumnStatistics::new()
1337            .with_row_count(10)
1338            .with_row_count(20);
1339        assert_eq!(stats.row_count, Some(20));
1340    }
1341
1342    #[test]
1343    fn column_statistics_eq() {
1344        let a = ColumnStatistics::new()
1345            .with_row_count(100)
1346            .with_null_count(5);
1347        let b = ColumnStatistics::new()
1348            .with_row_count(100)
1349            .with_null_count(5);
1350        let c = ColumnStatistics::new().with_row_count(99);
1351        assert_eq!(a, b);
1352        assert_ne!(a, c);
1353    }
1354
1355    // -----------------------------------------------------------------------
1356    // TableMetadata tests
1357    // -----------------------------------------------------------------------
1358
1359    #[test]
1360    fn table_metadata_new_and_accessors() {
1361        let meta = TableMetadata::new("events", make_schema());
1362        assert_eq!(meta.name(), "events");
1363        assert_eq!(meta.schema().field_count(), 2);
1364        assert!(meta.statistics().is_none());
1365    }
1366
1367    #[test]
1368    fn table_metadata_with_stats() {
1369        let stats = ColumnStatistics::new()
1370            .with_row_count(5000)
1371            .with_null_count(100);
1372        let meta = TableMetadata::new("clicks", make_schema()).with_stats(stats);
1373        assert_eq!(meta.name(), "clicks");
1374        let s = meta.statistics().unwrap();
1375        assert_eq!(s.row_count, Some(5000));
1376        assert_eq!(s.null_count, Some(100));
1377    }
1378
1379    #[test]
1380    fn table_metadata_into_string() {
1381        let meta = TableMetadata::new("test_table", make_schema());
1382        let dbg = format!("{meta:?}");
1383        assert!(dbg.contains("name: \"test_table\""));
1384        assert!(dbg.contains("stats: None"));
1385    }
1386
1387    // -----------------------------------------------------------------------
1388    // InMemorySchemaRegistry replace behavior
1389    // -----------------------------------------------------------------------
1390
1391    #[test]
1392    fn schema_registry_register_replaces_existing() {
1393        let mut registry = InMemorySchemaRegistry::new();
1394        let schema_a = TableSchema::new(vec![CatalogField::new("a", FieldType::Int32, false)]);
1395        let schema_b = TableSchema::new(vec![CatalogField::new("b", FieldType::Utf8, true)]);
1396
1397        registry.register_schema("my_schema", schema_a);
1398        registry.register_schema("my_schema", schema_b);
1399
1400        let retrieved = registry.get_schema("my_schema").unwrap();
1401        assert_eq!(retrieved.field_count(), 1);
1402        assert_eq!(
1403            retrieved.get_field("b").unwrap().field_type(),
1404            &FieldType::Utf8
1405        );
1406        assert!(retrieved.get_field("a").is_none());
1407    }
1408
1409    #[test]
1410    fn schema_registry_empty_names() {
1411        let registry = InMemorySchemaRegistry::new();
1412        assert!(registry.schema_names().is_empty());
1413    }
1414
1415    // -----------------------------------------------------------------------
1416    // TableSchema accessors
1417    // -----------------------------------------------------------------------
1418
1419    #[test]
1420    fn table_schema_empty() {
1421        let schema = TableSchema::empty();
1422        assert_eq!(schema.field_count(), 0);
1423        assert!(schema.get_field("anything").is_none());
1424    }
1425
1426    #[test]
1427    fn table_schema_get_field_found() {
1428        let schema = make_schema();
1429        let field = schema.get_field("name").unwrap();
1430        assert_eq!(field.name(), "name");
1431        assert_eq!(field.field_type(), &FieldType::Utf8);
1432        assert!(field.nullable());
1433    }
1434
1435    #[test]
1436    fn table_schema_get_field_not_found() {
1437        let schema = make_schema();
1438        assert!(schema.get_field("missing").is_none());
1439    }
1440
1441    // -----------------------------------------------------------------------
1442    // InMemoryCatalog duplicate registration
1443    // -----------------------------------------------------------------------
1444
1445    #[test]
1446    fn in_memory_catalog_duplicate_register_errors() {
1447        let mut catalog = InMemoryCatalog::new();
1448        catalog
1449            .register_table(TableMetadata::new("t", make_schema()))
1450            .unwrap();
1451        let err = catalog
1452            .register_table(TableMetadata::new("t", make_schema()))
1453            .unwrap_err();
1454        match err {
1455            CatalogError::TableAlreadyExists { name } => assert_eq!(name, "t"),
1456            other => panic!("expected TableAlreadyExists, got {other}"),
1457        }
1458    }
1459
1460    // -----------------------------------------------------------------------
1461    // CatalogResult type alias
1462    // -----------------------------------------------------------------------
1463
1464    #[test]
1465    fn catalog_result_ok() {
1466        let r: CatalogResult<i32> = Ok(42);
1467        assert_eq!(r.ok(), Some(42));
1468    }
1469
1470    #[test]
1471    fn catalog_result_err() {
1472        let r: CatalogResult<()> = Err(CatalogError::TableNotFound {
1473            name: "x".to_string(),
1474        });
1475        assert!(r.is_err());
1476    }
1477
1478    // -----------------------------------------------------------------------
1479    // Edge cases: empty catalog
1480    // -----------------------------------------------------------------------
1481
1482    #[test]
1483    fn empty_catalog_list_tables_returns_empty() {
1484        let catalog = InMemoryCatalog::new();
1485        assert!(catalog.list_tables().is_empty());
1486    }
1487
1488    #[test]
1489    fn empty_catalog_get_table_returns_not_found() {
1490        let catalog = InMemoryCatalog::new();
1491        let err = catalog.get_table("anything").err().unwrap();
1492        assert!(matches!(err, CatalogError::TableNotFound { .. }));
1493    }
1494
1495    #[test]
1496    fn empty_schema_registry_get_returns_not_found() {
1497        let registry = InMemorySchemaRegistry::new();
1498        assert!(registry.get_schema("x").is_err());
1499    }
1500
1501    #[test]
1502    fn empty_schema_schema_names_empty() {
1503        let registry = InMemorySchemaRegistry::new();
1504        assert!(registry.schema_names().is_empty());
1505    }
1506
1507    // -----------------------------------------------------------------------
1508    // Edge cases: special characters in names
1509    // -----------------------------------------------------------------------
1510
1511    #[test]
1512    fn table_name_with_special_characters() {
1513        let mut catalog = InMemoryCatalog::new();
1514        let meta = TableMetadata::new("table-with-dashes.dots_and_underscores", make_schema());
1515        catalog.register_table(meta).unwrap();
1516        let table = catalog
1517            .get_table("table-with-dashes.dots_and_underscores")
1518            .unwrap();
1519        assert_eq!(table.name(), "table-with-dashes.dots_and_underscores");
1520    }
1521
1522    #[test]
1523    fn table_name_with_unicode() {
1524        let mut catalog = InMemoryCatalog::new();
1525        let meta = TableMetadata::new("用户_table", make_schema());
1526        catalog.register_table(meta).unwrap();
1527        let table = catalog.get_table("用户_table").unwrap();
1528        assert_eq!(table.name(), "用户_table");
1529    }
1530
1531    #[test]
1532    fn table_name_with_spaces() {
1533        let mut catalog = InMemoryCatalog::new();
1534        let meta = TableMetadata::new("my table name", make_schema());
1535        catalog.register_table(meta).unwrap();
1536        let table = catalog.get_table("my table name").unwrap();
1537        assert_eq!(table.name(), "my table name");
1538    }
1539
1540    #[test]
1541    fn schema_name_with_special_characters() {
1542        let mut registry = InMemorySchemaRegistry::new();
1543        let schema = TableSchema::new(vec![CatalogField::new("col", FieldType::Int32, true)]);
1544        registry.register_schema("schema-with-dashes", schema);
1545        let retrieved = registry.get_schema("schema-with-dashes").unwrap();
1546        assert_eq!(retrieved.field_count(), 1);
1547    }
1548
1549    #[test]
1550    fn field_name_with_special_characters() {
1551        let f = CatalogField::new("field-with-dots_and@spaces", FieldType::Utf8, false);
1552        assert_eq!(f.name(), "field-with-dots_and@spaces");
1553        let arrow_f = f.to_arrow_field();
1554        assert_eq!(arrow_f.name(), "field-with-dots_and@spaces");
1555    }
1556
1557    // -----------------------------------------------------------------------
1558    // Edge cases: duplicate registration and overwrite
1559    // -----------------------------------------------------------------------
1560
1561    #[test]
1562    fn catalog_duplicate_different_table_errors() {
1563        let mut catalog = InMemoryCatalog::new();
1564        catalog
1565            .register_table(TableMetadata::new("t1", make_schema()))
1566            .unwrap();
1567        catalog
1568            .register_table(TableMetadata::new("t2", make_schema()))
1569            .unwrap();
1570        assert_eq!(catalog.list_tables().len(), 2);
1571    }
1572
1573    #[test]
1574    fn schema_registry_overwrite_preserves_single_entry() {
1575        let mut registry = InMemorySchemaRegistry::new();
1576        registry.register_schema("s", TableSchema::empty());
1577        registry.register_schema("s", make_schema());
1578        assert_eq!(registry.schema_names().len(), 1);
1579        assert_eq!(registry.get_schema("s").unwrap().field_count(), 2);
1580    }
1581
1582    // -----------------------------------------------------------------------
1583    // TableSchema: empty schema arrow conversion
1584    // -----------------------------------------------------------------------
1585
1586    #[test]
1587    fn empty_schema_to_arrow() {
1588        let schema = TableSchema::empty();
1589        let arrow_schema = schema.to_arrow_schema();
1590        assert_eq!(arrow_schema.fields().len(), 0);
1591    }
1592
1593    #[test]
1594    fn single_field_schema_to_arrow() {
1595        let schema = TableSchema::new(vec![CatalogField::new("only", FieldType::Float32, true)]);
1596        let arrow_schema = schema.to_arrow_schema();
1597        assert_eq!(arrow_schema.fields().len(), 1);
1598        let f = arrow_schema.field_with_name("only").unwrap();
1599        assert_eq!(f.data_type(), &arrow::datatypes::DataType::Float32);
1600        assert!(f.is_nullable());
1601    }
1602
1603    // -----------------------------------------------------------------------
1604    // FieldType: nested types to_arrow
1605    // -----------------------------------------------------------------------
1606
1607    #[test]
1608    fn field_type_list_of_list() {
1609        let inner = FieldType::List(Box::new(FieldType::Int32));
1610        let outer = FieldType::List(Box::new(inner));
1611        match outer.to_arrow() {
1612            arrow::datatypes::DataType::List(field) => match field.data_type() {
1613                arrow::datatypes::DataType::List(inner_field) => {
1614                    assert_eq!(inner_field.data_type(), &arrow::datatypes::DataType::Int32);
1615                }
1616                other => panic!("expected nested List, got {other:?}"),
1617            },
1618            other => panic!("expected outer List, got {other:?}"),
1619        }
1620    }
1621
1622    #[test]
1623    fn field_type_struct_nested_in_struct() {
1624        let inner_struct =
1625            FieldType::Struct(vec![CatalogField::new("a", FieldType::Boolean, true)]);
1626        let outer_struct = FieldType::Struct(vec![
1627            CatalogField::new("nested", inner_struct, false),
1628            CatalogField::new("simple", FieldType::Utf8, true),
1629        ]);
1630        match outer_struct.to_arrow() {
1631            arrow::datatypes::DataType::Struct(fields) => {
1632                assert_eq!(fields.len(), 2);
1633                match fields[0].data_type() {
1634                    arrow::datatypes::DataType::Struct(inner_fields) => {
1635                        assert_eq!(inner_fields.len(), 1);
1636                        assert_eq!(inner_fields[0].name(), "a");
1637                    }
1638                    other => panic!("expected inner Struct, got {other:?}"),
1639                }
1640                assert_eq!(fields[1].data_type(), &arrow::datatypes::DataType::Utf8);
1641            }
1642            other => panic!("expected Struct, got {other:?}"),
1643        }
1644    }
1645
1646    #[test]
1647    fn field_type_list_of_struct() {
1648        let list_type = FieldType::List(Box::new(FieldType::Struct(vec![
1649            CatalogField::new("x", FieldType::Int64, false),
1650            CatalogField::new("y", FieldType::Utf8, true),
1651        ])));
1652        match list_type.to_arrow() {
1653            arrow::datatypes::DataType::List(item_field) => match item_field.data_type() {
1654                arrow::datatypes::DataType::Struct(fields) => {
1655                    assert_eq!(fields.len(), 2);
1656                }
1657                other => panic!("expected inner Struct, got {other:?}"),
1658            },
1659            other => panic!("expected List, got {other:?}"),
1660        }
1661    }
1662
1663    #[test]
1664    fn field_type_empty_struct() {
1665        let empty_struct = FieldType::Struct(vec![]);
1666        match empty_struct.to_arrow() {
1667            arrow::datatypes::DataType::Struct(fields) => {
1668                assert_eq!(fields.len(), 0);
1669            }
1670            other => panic!("expected Struct, got {other:?}"),
1671        }
1672    }
1673
1674    #[test]
1675    fn field_type_list_of_binary() {
1676        let list_type = FieldType::List(Box::new(FieldType::Binary));
1677        match list_type.to_arrow() {
1678            arrow::datatypes::DataType::List(field) => {
1679                assert_eq!(field.data_type(), &arrow::datatypes::DataType::Binary);
1680            }
1681            other => panic!("expected List, got {other:?}"),
1682        }
1683    }
1684
1685    // -----------------------------------------------------------------------
1686    // CatalogField: Clone and PartialEq
1687    // -----------------------------------------------------------------------
1688
1689    #[test]
1690    fn catalog_field_clone_eq() {
1691        let f1 = CatalogField::new("col", FieldType::Int32, true);
1692        let f2 = f1.clone();
1693        assert_eq!(f1, f2);
1694    }
1695
1696    #[test]
1697    fn catalog_field_ne_name() {
1698        let f1 = CatalogField::new("a", FieldType::Int32, true);
1699        let f2 = CatalogField::new("b", FieldType::Int32, true);
1700        assert_ne!(f1, f2);
1701    }
1702
1703    #[test]
1704    fn catalog_field_ne_type() {
1705        let f1 = CatalogField::new("a", FieldType::Int32, true);
1706        let f2 = CatalogField::new("a", FieldType::Utf8, true);
1707        assert_ne!(f1, f2);
1708    }
1709
1710    #[test]
1711    fn catalog_field_ne_nullable() {
1712        let f1 = CatalogField::new("a", FieldType::Int32, true);
1713        let f2 = CatalogField::new("a", FieldType::Int32, false);
1714        assert_ne!(f1, f2);
1715    }
1716
1717    // -----------------------------------------------------------------------
1718    // TableSchema: Clone and PartialEq
1719    // -----------------------------------------------------------------------
1720
1721    #[test]
1722    fn table_schema_clone_eq() {
1723        let s1 = make_schema();
1724        let s2 = s1.clone();
1725        assert_eq!(s1, s2);
1726    }
1727
1728    #[test]
1729    fn table_schema_ne_different_fields() {
1730        let s1 = TableSchema::new(vec![CatalogField::new("a", FieldType::Int32, false)]);
1731        let s2 = TableSchema::new(vec![CatalogField::new("b", FieldType::Int32, false)]);
1732        assert_ne!(s1, s2);
1733    }
1734
1735    // -----------------------------------------------------------------------
1736    // TableMetadata: Clone
1737    // -----------------------------------------------------------------------
1738
1739    #[test]
1740    fn table_metadata_clone() {
1741        let meta = TableMetadata::new("t", make_schema())
1742            .with_stats(ColumnStatistics::new().with_row_count(100));
1743        let cloned = meta.clone();
1744        assert_eq!(cloned.name(), "t");
1745        assert_eq!(cloned.statistics().unwrap().row_count, Some(100));
1746    }
1747
1748    // -----------------------------------------------------------------------
1749    // InMemoryCatalog: register_table_with_batches
1750    // -----------------------------------------------------------------------
1751
1752    #[test]
1753    fn register_table_with_batches_stores_data() {
1754        let mut catalog = InMemoryCatalog::new();
1755        let schema = TableSchema::new(vec![CatalogField::new("id", FieldType::Int64, false)]);
1756        let arrow_schema = std::sync::Arc::new(arrow::datatypes::Schema::new(vec![
1757            arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
1758        ]));
1759        let batch = arrow::record_batch::RecordBatch::try_new(
1760            arrow_schema,
1761            vec![std::sync::Arc::new(arrow::array::Int64Array::from(vec![
1762                1, 2, 3,
1763            ]))],
1764        )
1765        .unwrap();
1766        catalog
1767            .register_table_with_batches(TableMetadata::new("data", schema), vec![batch])
1768            .unwrap();
1769        assert!(catalog.table_batches("data").is_some());
1770        assert_eq!(catalog.table_batches("data").unwrap().len(), 1);
1771        assert_eq!(catalog.table_batches("data").unwrap()[0].num_rows(), 3);
1772    }
1773
1774    #[test]
1775    fn register_table_with_empty_batches_no_data() {
1776        let mut catalog = InMemoryCatalog::new();
1777        let schema = TableSchema::new(vec![CatalogField::new("id", FieldType::Int32, false)]);
1778        catalog
1779            .register_table_with_batches(TableMetadata::new("empty", schema), vec![])
1780            .unwrap();
1781        assert!(catalog.table_batches("empty").is_none());
1782    }
1783
1784    #[test]
1785    fn table_batches_nonexistent_table() {
1786        let catalog = InMemoryCatalog::new();
1787        assert!(catalog.table_batches("nope").is_none());
1788    }
1789
1790    // -----------------------------------------------------------------------
1791    // InMemoryCatalog: Default
1792    // -----------------------------------------------------------------------
1793
1794    #[test]
1795    fn in_memory_catalog_default() {
1796        let catalog = InMemoryCatalog::default();
1797        assert!(catalog.list_tables().is_empty());
1798    }
1799
1800    // -----------------------------------------------------------------------
1801    // InMemorySchemaRegistry: Default
1802    // -----------------------------------------------------------------------
1803
1804    #[test]
1805    fn in_memory_schema_registry_default() {
1806        let registry = InMemorySchemaRegistry::default();
1807        assert!(registry.schema_names().is_empty());
1808    }
1809
1810    // -----------------------------------------------------------------------
1811    // DataFusion bridge: SQL with empty table
1812    // -----------------------------------------------------------------------
1813
1814    #[tokio::test]
1815    async fn datafusion_bridge_empty_table_query() {
1816        use datafusion::prelude::SessionContext;
1817
1818        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
1819        {
1820            let mut cat = catalog.write().unwrap();
1821            let schema = TableSchema::new(vec![CatalogField::new("id", FieldType::Int64, false)]);
1822            cat.register_table(TableMetadata::new("empty_table", schema))
1823                .unwrap();
1824        }
1825        let ctx = SessionContext::new();
1826        ctx.register_catalog(
1827            "krishiv",
1828            std::sync::Arc::new(super::datafusion_bridge::DataFusionCatalogBridge::new(
1829                catalog,
1830            )),
1831        );
1832        let df = ctx
1833            .sql("SELECT * FROM krishiv.public.empty_table")
1834            .await
1835            .unwrap();
1836        let batches = df.collect().await.unwrap();
1837        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1838        assert_eq!(rows, 0);
1839    }
1840
1841    #[tokio::test]
1842    async fn datafusion_bridge_sql_filter() {
1843        use std::sync::Arc;
1844
1845        use arrow::array::Int64Array;
1846        use arrow::datatypes::{DataType, Field, Schema};
1847        use arrow::record_batch::RecordBatch;
1848        use datafusion::prelude::SessionContext;
1849
1850        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
1851        let schema = TableSchema::new(vec![
1852            CatalogField::new("id", FieldType::Int64, false),
1853            CatalogField::new("val", FieldType::Int64, false),
1854        ]);
1855        let arrow_schema = Arc::new(Schema::new(vec![
1856            Field::new("id", DataType::Int64, false),
1857            Field::new("val", DataType::Int64, false),
1858        ]));
1859        let batch = RecordBatch::try_new(
1860            arrow_schema,
1861            vec![
1862                Arc::new(Int64Array::from(vec![1, 2, 3])),
1863                Arc::new(Int64Array::from(vec![10, 20, 30])),
1864            ],
1865        )
1866        .unwrap();
1867        catalog
1868            .write()
1869            .unwrap()
1870            .register_table_with_batches(TableMetadata::new("nums", schema), vec![batch])
1871            .unwrap();
1872
1873        let ctx = SessionContext::new();
1874        ctx.register_catalog(
1875            "krishiv",
1876            Arc::new(super::datafusion_bridge::DataFusionCatalogBridge::new(
1877                catalog,
1878            )),
1879        );
1880        let df = ctx
1881            .sql("SELECT id FROM krishiv.public.nums WHERE val > 15")
1882            .await
1883            .unwrap();
1884        let batches = df.collect().await.unwrap();
1885        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1886        assert_eq!(rows, 2);
1887    }
1888
1889    #[tokio::test]
1890    async fn datafusion_bridge_sql_count_aggregate() {
1891        use std::sync::Arc;
1892
1893        use arrow::datatypes::{DataType, Field, Schema};
1894        use arrow::record_batch::RecordBatch;
1895        use datafusion::prelude::SessionContext;
1896
1897        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
1898        let schema = TableSchema::new(vec![CatalogField::new("x", FieldType::Int32, false)]);
1899        let arrow_schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
1900        let batch = RecordBatch::try_new(
1901            arrow_schema,
1902            vec![Arc::new(arrow::array::Int32Array::from(vec![
1903                1, 2, 3, 4, 5,
1904            ]))],
1905        )
1906        .unwrap();
1907        catalog
1908            .write()
1909            .unwrap()
1910            .register_table_with_batches(TableMetadata::new("agg", schema), vec![batch])
1911            .unwrap();
1912
1913        let ctx = SessionContext::new();
1914        ctx.register_catalog(
1915            "krishiv",
1916            Arc::new(super::datafusion_bridge::DataFusionCatalogBridge::new(
1917                catalog,
1918            )),
1919        );
1920        let df = ctx
1921            .sql("SELECT COUNT(*) AS cnt FROM krishiv.public.agg")
1922            .await
1923            .unwrap();
1924        let batches = df.collect().await.unwrap();
1925        assert_eq!(batches.len(), 1);
1926        assert_eq!(batches[0].num_rows(), 1);
1927    }
1928
1929    // -----------------------------------------------------------------------
1930    // DataFusion bridge: multiple tables
1931    // -----------------------------------------------------------------------
1932
1933    #[tokio::test]
1934    async fn datafusion_bridge_multiple_tables() {
1935        use std::sync::Arc;
1936
1937        use arrow::array::Int64Array;
1938        use arrow::datatypes::{DataType, Field, Schema};
1939        use arrow::record_batch::RecordBatch;
1940        use datafusion::prelude::SessionContext;
1941
1942        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
1943        let schema = TableSchema::new(vec![CatalogField::new("id", FieldType::Int64, false)]);
1944        let arrow_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
1945        let batch =
1946            RecordBatch::try_new(arrow_schema, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap();
1947        {
1948            let mut cat = catalog.write().unwrap();
1949            cat.register_table_with_batches(
1950                TableMetadata::new("t1", schema.clone()),
1951                vec![batch.clone()],
1952            )
1953            .unwrap();
1954            cat.register_table_with_batches(TableMetadata::new("t2", schema), vec![batch])
1955                .unwrap();
1956        }
1957
1958        let ctx = SessionContext::new();
1959        ctx.register_catalog(
1960            "krishiv",
1961            Arc::new(super::datafusion_bridge::DataFusionCatalogBridge::new(
1962                catalog,
1963            )),
1964        );
1965        let df1 = ctx.sql("SELECT * FROM krishiv.public.t1").await.unwrap();
1966        let batches1 = df1.collect().await.unwrap();
1967        assert_eq!(batches1.len(), 1);
1968        assert_eq!(batches1[0].num_rows(), 1);
1969
1970        let df2 = ctx.sql("SELECT * FROM krishiv.public.t2").await.unwrap();
1971        let batches2 = df2.collect().await.unwrap();
1972        assert_eq!(batches2.len(), 1);
1973        assert_eq!(batches2[0].num_rows(), 1);
1974    }
1975
1976    // -----------------------------------------------------------------------
1977    // DataFusion bridge: schema() returns None for non-public
1978    // -----------------------------------------------------------------------
1979
1980    #[test]
1981    fn datafusion_bridge_custom_schema_name_returns_none() {
1982        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
1983
1984        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
1985        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
1986        assert!(bridge.schema("custom").is_none());
1987        assert!(bridge.schema("public").is_some());
1988    }
1989
1990    // -----------------------------------------------------------------------
1991    // DataFusion bridge: as_any returns self
1992    // -----------------------------------------------------------------------
1993
1994    // -----------------------------------------------------------------------
1995    // DataFusion bridge: schema_names returns single "public"
1996    // -----------------------------------------------------------------------
1997
1998    #[test]
1999    fn datafusion_bridge_only_public_schema() {
2000        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
2001
2002        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2003        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2004        let names = bridge.schema_names();
2005        assert_eq!(names.len(), 1);
2006        assert_eq!(names[0], "public");
2007    }
2008
2009    // -----------------------------------------------------------------------
2010    // FieldType Display: empty struct and empty list
2011    // -----------------------------------------------------------------------
2012
2013    #[test]
2014    fn field_type_display_empty_struct() {
2015        let s = FieldType::Struct(vec![]);
2016        assert_eq!(s.to_string(), "Struct(0 fields)");
2017    }
2018
2019    #[test]
2020    fn field_type_display_nested_list() {
2021        let inner = FieldType::List(Box::new(FieldType::Int32));
2022        let outer = FieldType::List(Box::new(inner));
2023        assert_eq!(outer.to_string(), "List<List<Int32>>");
2024    }
2025
2026    // -----------------------------------------------------------------------
2027    // ColumnStatistics: eq with None fields
2028    // -----------------------------------------------------------------------
2029
2030    #[test]
2031    fn column_statistics_eq_all_none() {
2032        let a = ColumnStatistics::new();
2033        let b = ColumnStatistics::new();
2034        assert_eq!(a, b);
2035    }
2036
2037    #[test]
2038    fn column_statistics_ne_different_min() {
2039        let a = ColumnStatistics::new().with_min("aaa");
2040        let b = ColumnStatistics::new().with_min("zzz");
2041        assert_ne!(a, b);
2042    }
2043
2044    // -----------------------------------------------------------------------
2045    // TableSchema: get_field with multiple fields
2046    // -----------------------------------------------------------------------
2047
2048    #[test]
2049    fn table_schema_get_field_last() {
2050        let schema = TableSchema::new(vec![
2051            CatalogField::new("a", FieldType::Int32, false),
2052            CatalogField::new("b", FieldType::Utf8, true),
2053            CatalogField::new("c", FieldType::Float64, false),
2054        ]);
2055        let field = schema.get_field("c").unwrap();
2056        assert_eq!(field.field_type(), &FieldType::Float64);
2057    }
2058
2059    #[test]
2060    fn table_schema_get_field_middle() {
2061        let schema = TableSchema::new(vec![
2062            CatalogField::new("a", FieldType::Int32, false),
2063            CatalogField::new("b", FieldType::Utf8, true),
2064            CatalogField::new("c", FieldType::Float64, false),
2065        ]);
2066        let field = schema.get_field("b").unwrap();
2067        assert_eq!(field.field_type(), &FieldType::Utf8);
2068        assert!(field.nullable());
2069    }
2070
2071    // -----------------------------------------------------------------------
2072    // CatalogError: all variants as std::error::Error
2073    // -----------------------------------------------------------------------
2074
2075    #[test]
2076    fn all_catalog_errors_are_std_error() {
2077        let errors: Vec<CatalogError> = vec![
2078            CatalogError::TableNotFound { name: "a".into() },
2079            CatalogError::TableAlreadyExists { name: "b".into() },
2080            CatalogError::SchemaNotFound { name: "c".into() },
2081            CatalogError::InvalidSchema {
2082                message: "d".into(),
2083            },
2084            CatalogError::InvalidConfiguration {
2085                message: "bad URL".into(),
2086            },
2087            CatalogError::Transport {
2088                operation: "load table".into(),
2089                message: "timed out".into(),
2090            },
2091            CatalogError::Http {
2092                status: 500,
2093                message: "e".into(),
2094            },
2095            CatalogError::InvalidResponse {
2096                operation: "list tables".into(),
2097                message: "missing identifiers".into(),
2098            },
2099            CatalogError::ResponseTooLarge {
2100                operation: "load table".into(),
2101                limit_bytes: 1024,
2102            },
2103            CatalogError::UnsupportedOperation {
2104                operation: "committing a table".into(),
2105            },
2106        ];
2107        for err in errors {
2108            let e: &dyn std::error::Error = &err;
2109            let _ = e.to_string();
2110            assert!(e.source().is_none());
2111        }
2112    }
2113
2114    // -----------------------------------------------------------------------
2115    // InMemoryCatalog: register many tables
2116    // -----------------------------------------------------------------------
2117
2118    #[test]
2119    fn in_memory_catalog_many_tables() {
2120        let mut catalog = InMemoryCatalog::new();
2121        for i in 0..100 {
2122            catalog
2123                .register_table(TableMetadata::new(format!("table_{i:03}"), make_schema()))
2124                .unwrap();
2125        }
2126        assert_eq!(catalog.list_tables().len(), 100);
2127        assert!(catalog.get_table("table_000").is_ok());
2128        assert!(catalog.get_table("table_099").is_ok());
2129        assert!(catalog.get_table("table_100").is_err());
2130    }
2131
2132    // -----------------------------------------------------------------------
2133    // DataFusion bridge: table() returns None for unknown
2134    // -----------------------------------------------------------------------
2135
2136    #[tokio::test]
2137    async fn datafusion_bridge_table_unknown() {
2138        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
2139        use std::sync::Arc;
2140
2141        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2142        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2143        let schema_provider = bridge.schema("public").unwrap();
2144        let result = schema_provider.table("nonexistent").await.unwrap();
2145        assert!(result.is_none());
2146    }
2147
2148    // -----------------------------------------------------------------------
2149    // DataFusion bridge: table() returns Some for known table
2150    // -----------------------------------------------------------------------
2151
2152    #[tokio::test]
2153    async fn datafusion_bridge_table_known() {
2154        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
2155        use std::sync::Arc;
2156
2157        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2158        {
2159            let mut cat = catalog.write().unwrap();
2160            cat.register_table(TableMetadata::new("mytable", make_schema()))
2161                .unwrap();
2162        }
2163        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2164        let schema_provider = bridge.schema("public").unwrap();
2165        let result = schema_provider.table("mytable").await.unwrap();
2166        assert!(result.is_some());
2167    }
2168
2169    // -----------------------------------------------------------------------
2170    // DataFusion bridge: table_names lists registered tables
2171    // -----------------------------------------------------------------------
2172
2173    #[tokio::test]
2174    async fn datafusion_bridge_table_names() {
2175        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
2176        use std::sync::Arc;
2177
2178        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2179        {
2180            let mut cat = catalog.write().unwrap();
2181            cat.register_table(TableMetadata::new("alpha", make_schema()))
2182                .unwrap();
2183            cat.register_table(TableMetadata::new("beta", make_schema()))
2184                .unwrap();
2185        }
2186        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2187        let schema_provider = bridge.schema("public").unwrap();
2188        let mut names = schema_provider.table_names();
2189        names.sort();
2190        assert_eq!(names, vec!["alpha", "beta"]);
2191    }
2192
2193    // -----------------------------------------------------------------------
2194    // DataFusion bridge: empty catalog table_names
2195    // -----------------------------------------------------------------------
2196
2197    #[tokio::test]
2198    async fn datafusion_bridge_empty_table_names() {
2199        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
2200        use std::sync::Arc;
2201
2202        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2203        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2204        let schema_provider = bridge.schema("public").unwrap();
2205        let names = schema_provider.table_names();
2206        assert!(names.is_empty());
2207    }
2208
2209    // -----------------------------------------------------------------------
2210    // DataFusion bridge: table_exist with multiple tables
2211    // -----------------------------------------------------------------------
2212
2213    #[test]
2214    fn datafusion_bridge_table_exist_multiple() {
2215        use datafusion::catalog::CatalogProvider as DfCatalogProvider;
2216
2217        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2218        {
2219            let mut cat = catalog.write().unwrap();
2220            cat.register_table(TableMetadata::new("a", make_schema()))
2221                .unwrap();
2222            cat.register_table(TableMetadata::new("b", make_schema()))
2223                .unwrap();
2224        }
2225        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2226        let sp = bridge.schema("public").unwrap();
2227        assert!(sp.table_exist("a"));
2228        assert!(sp.table_exist("b"));
2229        assert!(!sp.table_exist("c"));
2230    }
2231
2232    // -----------------------------------------------------------------------
2233    // DataFusion bridge: debug format
2234    // -----------------------------------------------------------------------
2235
2236    #[test]
2237    fn datafusion_bridge_debug() {
2238        let catalog = std::sync::Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2239        let bridge = super::datafusion_bridge::DataFusionCatalogBridge::new(catalog);
2240        let dbg = format!("{bridge:?}");
2241        assert!(dbg.contains("DataFusionCatalogBridge"));
2242    }
2243
2244    // -----------------------------------------------------------------------
2245    // InMemoryCatalog: register_table_with_batches error on duplicate
2246    // -----------------------------------------------------------------------
2247
2248    #[test]
2249    fn register_table_with_batches_duplicate_errors() {
2250        let mut catalog = InMemoryCatalog::new();
2251        let schema = TableSchema::new(vec![CatalogField::new("id", FieldType::Int32, false)]);
2252        let arrow_schema = std::sync::Arc::new(arrow::datatypes::Schema::new(vec![
2253            arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
2254        ]));
2255        let batch = arrow::record_batch::RecordBatch::try_new(
2256            arrow_schema,
2257            vec![std::sync::Arc::new(arrow::array::Int32Array::from(vec![1]))],
2258        )
2259        .unwrap();
2260        catalog
2261            .register_table_with_batches(TableMetadata::new("t", schema.clone()), vec![batch])
2262            .unwrap();
2263        let err = catalog
2264            .register_table_with_batches(TableMetadata::new("t", schema), vec![])
2265            .unwrap_err();
2266        assert!(matches!(err, CatalogError::TableAlreadyExists { .. }));
2267    }
2268
2269    // -----------------------------------------------------------------------
2270    // DataFusion bridge: SQL with multiple columns and types
2271    // -----------------------------------------------------------------------
2272
2273    #[tokio::test]
2274    async fn datafusion_bridge_sql_multiple_columns() {
2275        use std::sync::Arc;
2276
2277        use arrow::array::{Int32Array, StringArray};
2278        use arrow::datatypes::{DataType, Field, Schema};
2279        use arrow::record_batch::RecordBatch;
2280        use datafusion::prelude::SessionContext;
2281
2282        let catalog = Arc::new(std::sync::RwLock::new(InMemoryCatalog::new()));
2283        let schema = TableSchema::new(vec![
2284            CatalogField::new("id", FieldType::Int32, false),
2285            CatalogField::new("name", FieldType::Utf8, true),
2286        ]);
2287        let arrow_schema = Arc::new(Schema::new(vec![
2288            Field::new("id", DataType::Int32, false),
2289            Field::new("name", DataType::Utf8, true),
2290        ]));
2291        let batch = RecordBatch::try_new(
2292            arrow_schema,
2293            vec![
2294                Arc::new(Int32Array::from(vec![1, 2, 3])),
2295                Arc::new(StringArray::from(vec!["a", "b", "c"])),
2296            ],
2297        )
2298        .unwrap();
2299        catalog
2300            .write()
2301            .unwrap()
2302            .register_table_with_batches(TableMetadata::new("mixed", schema), vec![batch])
2303            .unwrap();
2304
2305        let ctx = SessionContext::new();
2306        ctx.register_catalog(
2307            "krishiv",
2308            Arc::new(super::datafusion_bridge::DataFusionCatalogBridge::new(
2309                catalog,
2310            )),
2311        );
2312        let df = ctx
2313            .sql("SELECT name FROM krishiv.public.mixed WHERE id > 1")
2314            .await
2315            .unwrap();
2316        let batches = df.collect().await.unwrap();
2317        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
2318        assert_eq!(rows, 2);
2319    }
2320
2321    // -----------------------------------------------------------------------
2322    // CatalogError: Display for all variants via format
2323    // -----------------------------------------------------------------------
2324
2325    #[test]
2326    fn catalog_error_display_all_variants() {
2327        let cases: Vec<(CatalogError, &str)> = vec![
2328            (
2329                CatalogError::TableNotFound {
2330                    name: "x".to_string(),
2331                },
2332                "table not found: 'x'",
2333            ),
2334            (
2335                CatalogError::TableAlreadyExists {
2336                    name: "y".to_string(),
2337                },
2338                "table already exists: 'y'",
2339            ),
2340            (
2341                CatalogError::SchemaNotFound {
2342                    name: "z".to_string(),
2343                },
2344                "schema not found: 'z'",
2345            ),
2346            (
2347                CatalogError::InvalidSchema {
2348                    message: "bad".to_string(),
2349                },
2350                "invalid schema: bad",
2351            ),
2352            (
2353                CatalogError::Http {
2354                    status: 403,
2355                    message: "forbidden".to_string(),
2356                },
2357                "HTTP error 403: forbidden",
2358            ),
2359            (
2360                CatalogError::InvalidConfiguration {
2361                    message: "bad URL".to_string(),
2362                },
2363                "invalid catalog configuration: bad URL",
2364            ),
2365            (
2366                CatalogError::Transport {
2367                    operation: "list tables".to_string(),
2368                    message: "timed out".to_string(),
2369                },
2370                "catalog transport error during list tables: timed out",
2371            ),
2372            (
2373                CatalogError::InvalidResponse {
2374                    operation: "load table".to_string(),
2375                    message: "missing metadata".to_string(),
2376                },
2377                "invalid catalog response during load table: missing metadata",
2378            ),
2379            (
2380                CatalogError::ResponseTooLarge {
2381                    operation: "load table".to_string(),
2382                    limit_bytes: 4096,
2383                },
2384                "catalog response during load table exceeded 4096 bytes",
2385            ),
2386            (
2387                CatalogError::UnsupportedOperation {
2388                    operation: "committing a table".to_string(),
2389                },
2390                "catalog server does not support committing a table",
2391            ),
2392        ];
2393        for (err, expected) in cases {
2394            assert_eq!(err.to_string(), expected);
2395        }
2396    }
2397}