Export and query

"""Export a Table to CSV, run a raw SQL SELECT, and ask a natural-language
question over the rows.

export() returns CSV/Excel bytes inline (or an s3_key when the server
streams to object storage). raw_query() runs SQL. nl_query() runs
server-side text-to-SQL and returns {answer, dataframes}.

import_file() has the same long-running shape as run_and_wait — see
the reference for that path.
"""

from vectorshift.table import (
    ColumnSpec,
    ExportFormat,
    NumberFormat,
    NumberKind,
    StringFormat,
    Table,
)

t = Table.new(
    name="SDK Tables Export Demo",
    columns=[
        ColumnSpec(name="vendor", format=StringFormat()),
        ColumnSpec(name="amount", format=NumberFormat(number_kind=NumberKind.FLOAT)),
    ],
)
t.insert_rows(
    [
        {"vendor": "Acme", "amount": 1200.5},
        {"vendor": "Beta", "amount": 49.9},
        {"vendor": "Gamma", "amount": 980.0},
    ]
)
print(f"1. Set up table id={t.id} with 3 rows.")

# --- Export to CSV ---
out = t.export(format=ExportFormat.CSV)
print(
    f"2. Exported -> filename={out['filename']!r} format={out['format']} bytes={len(out['content'])}"
)
print(f"     first 200 bytes: {out['content'][:200]!r}")

# --- raw_query: SELECT direct SQL ---
try:
    rows = t.raw_query("SELECT vendor FROM table_data LIMIT 5")
    print(f"3. raw_query returned {len(rows)} row(s).")
    for r in rows[:3]:
        print(f"     - row_id={r['row_id']} values={r['values']}")
except Exception as exc:
    print(f"3. raw_query error (engine surface): {type(exc).__name__}: {exc}")

# --- nl_query: server-side text-to-SQL ---
try:
    result = t.nl_query("how many rows are in this table?")
    print(f"4. nl_query answer: {result['answer']!r}")
    print(f"     dataframes returned: {len(result['dataframes'])}")
except Exception as exc:
    print(f"4. nl_query error: {type(exc).__name__}: {exc}")

t.delete()
print("5. Deleted Table.\n\nDone.")

1. Set up table id=... with 3 rows.
2. Exported -> filename='SDK Tables Export Demo.csv' format=ExportFormat.CSV bytes=...
     first 200 bytes: b'vendor,amount\nAcme,1200.5\nBeta,49.9\nGamma,980.0\n'
3. raw_query returned 3 row(s).
     - row_id=... values={'vendor': 'Acme'}
     ...
4. nl_query answer: '3'
     dataframes returned: 1
5. Deleted Table.

Done.