Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions python/cocoindex/cli.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import asyncio
import click
import datetime

from rich.console import Console
from rich.table import Table

from . import flow, lib, setting
from .setup import sync_setup, drop_setup, flow_names_with_setup, apply_setup_changes
Expand Down Expand Up @@ -56,11 +58,26 @@ def ls(show_all: bool):
@click.option("--color/--no-color", default=True)
def show(flow_name: str | None, color: bool):
"""
Show the flow spec in a readable format with colored output.
Show the flow spec in a readable format with colored output,
including the schema.
"""
fl = _flow_by_name(flow_name)
flow = _flow_by_name(flow_name)
console = Console(no_color=not color)
console.print(fl._render_text())
console.print(flow._render_text())

table = Table(
title=f"Schema for Flow: {flow.name}",
show_header=True,
header_style="bold magenta"
)
table.add_column("Field", style="cyan")
table.add_column("Type", style="green")
table.add_column("Attributes", style="yellow")

for field_name, field_type, attr_str in flow._render_schema():
table.add_row(field_name, field_type, attr_str)

console.print(table)

@cli.command()
def setup():
Expand Down
3 changes: 3 additions & 0 deletions python/cocoindex/flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,9 @@ def _render_text(self) -> Text:
return self._format_flow(flow_dict)
except json.JSONDecodeError:
return Text(flow_spec_str)

def _render_schema(self) -> list[tuple[str, str, str]]:
return self._lazy_engine_flow().get_schema()

def __str__(self):
return str(self._render_text())
Expand Down
65 changes: 65 additions & 0 deletions src/py/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::prelude::*;

use crate::base::schema::{BasicValueType, FieldSchema, ValueType};
use crate::base::spec::VectorSimilarityMetric;
use crate::execution::query;
use crate::lib_context::{clear_lib_context, get_auth_registry, init_lib_context};
Expand All @@ -13,6 +14,7 @@ use pyo3::{exceptions::PyException, prelude::*};
use pyo3_async_runtimes::tokio::future_into_py;
use std::collections::btree_map;
use std::fmt::Write;
use std::sync::Arc;

mod convert;
pub use convert::*;
Expand Down Expand Up @@ -193,6 +195,69 @@ impl Flow {
Ok(())
})
}

pub fn get_schema(&self) -> Vec<(String, String, String)> {
let schema = &self.0.flow.data_schema;
let mut result = Vec::new();

fn process_fields(
fields: &[FieldSchema],
prefix: &str,
result: &mut Vec<(String, String, String)>,
) {
for field in fields {
let field_name = format!("{}{}", prefix, field.name);

let mut field_type = match &field.value_type.typ {
ValueType::Basic(basic) => match basic {
BasicValueType::Vector(v) => {
let dim = v.dimension.map_or("*".to_string(), |d| d.to_string());
let elem = match *v.element_type {
BasicValueType::Float32 => "Float32",
BasicValueType::Float64 => "Float64",
_ => "Unknown",
};
format!("Vector[{}, {}]", dim, elem)
}
other => format!("{:?}", other),
},
ValueType::Table(t) => format!("{:?}", t.kind),
ValueType::Struct(_) => "Struct".to_string(),
};

if field.value_type.nullable {
field_type.push('?');
}

let attr_str = if field.value_type.attrs.is_empty() {
String::new()
} else {
field
.value_type
.attrs
.keys()
.map(|k| k.to_string())
.collect::<Vec<_>>()
.join(", ")
};

result.push((field_name.clone(), field_type, attr_str));

match &field.value_type.typ {
ValueType::Struct(s) => {
process_fields(&s.fields, &format!("{}.", field_name), result);
}
ValueType::Table(t) => {
process_fields(&t.row.fields, &format!("{}.", field_name), result);
}
ValueType::Basic(_) => {}
}
}
}

process_fields(&schema.schema.fields, "", &mut result);
result
}
}

#[pyclass]
Expand Down