re_datafusion 0.36.0

use std::collections::HashSet;
use std::pin::Pin;
use std::sync::Arc;

use crate::IntoDfError as _;
use crate::PendingTableQueryAnalytics;
use crate::analytics::QueryErrorKind;
use crate::batch_coalescer::coalesce_exec::SizedCoalesceBatchesExec;
use crate::batch_coalescer::coalescer::CoalescerOptions;
use arrow::array::{Array as _, RecordBatch, RecordBatchOptions};
use arrow::datatypes::SchemaRef;
use async_trait::async_trait;
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::not_impl_err;
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
use datafusion::logical_expr::TableProviderFilterPushDown;
use datafusion::logical_expr::dml::InsertOp;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec};
use datafusion::prelude::Expr;
use futures_util::StreamExt as _;
use re_redap_client::{ApiResponseStream, ApiResult};
use tokio_stream::Stream;

/// The projection, filters, and row limit offered to a scan.
///
/// Threaded from [`TableProvider::scan`] down to [`GrpcStreamToTable::send_streaming_request`]
/// (so an implementor can build a narrower request) and [`GrpcStreamToTable::process_response`].
/// Cheap to clone — the projection and filters are shared.
#[derive(Debug, Clone, Default)]
pub struct ScanParams {
    /// Unique column names required from the source, in scan order.
    ///
    /// `None` requests every column. `Some([])` is a zero-column DataFusion projection; protocols
    /// where an empty projection means "all columns" must request all columns and rely on the
    /// stream normalizer to remove them while preserving the row count.
    pub projected_columns: Option<Arc<[String]>>,

    /// Filters offered by DataFusion, combined with a logical AND.
    ///
    /// No filters = everything passes.
    pub filters: Arc<[Expr]>,

    /// Maximum number of rows the scan will consume, if bounded.
    ///
    /// Plumbed through so implementors can push it into their request once the backends grow
    /// server-side limit support; the coalescer already enforces it downstream in the meantime.
    #[expect(dead_code, reason = "plumbed for future server-side limit pushdown")]
    pub limit: Option<usize>,
}

#[async_trait]
pub trait GrpcStreamToTable:
    std::fmt::Debug + 'static + Send + Sync + Clone + std::marker::Unpin
{
    type GrpcStreamData;

    async fn fetch_schema(&mut self) -> ApiResult<SchemaRef>;

    fn process_response(
        &mut self,
        response: Self::GrpcStreamData,
        params: &ScanParams,
    ) -> ApiResult<RecordBatch>;

    async fn send_streaming_request(
        &mut self,
        params: &ScanParams,
    ) -> ApiResult<ApiResponseStream<Self::GrpcStreamData>>;

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
        Ok(vec![
            TableProviderFilterPushDown::Unsupported;
            filters.len()
        ])
    }

    async fn insert_into(
        &self,
        _state: &dyn Session,
        _input: Arc<dyn ExecutionPlan>,
        _insert_op: InsertOp,
    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
        not_impl_err!("Insert into not implemented for this table")
    }

    /// Optional analytics hook called once per `scan()`. Implementors that
    /// represent user-visible table scans (currently only
    /// `TableEntryTableProvider`) return a tracker that accumulates per-batch
    /// stats and emits an OTLP span on drop.
    fn begin_scan_analytics(
        &self,
        _schema: &SchemaRef,
        _projection: Option<&Vec<usize>>,
        _limit: Option<usize>,
    ) -> Option<PendingTableQueryAnalytics> {
        None
    }
}

#[derive(Debug)]
pub struct GrpcStreamProvider<T: GrpcStreamToTable> {
    schema: SchemaRef,
    client: T,
}

impl<T: GrpcStreamToTable> GrpcStreamProvider<T> {
    pub async fn prepare(mut client: T) -> Result<Arc<Self>, DataFusionError> {
        let schema = client
            .fetch_schema()
            .await
            .map_err(|err| err.into_df_error())?;
        Ok(Arc::new(Self { schema, client }))
    }
}

#[async_trait]
impl<T> TableProvider for GrpcStreamProvider<T>
where
    T: GrpcStreamToTable + Send + 'static,
    T::GrpcStreamData: Send + 'static,
{
    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    fn table_type(&self) -> datafusion::datasource::TableType {
        datafusion::datasource::TableType::Base
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        projection: Option<&Vec<usize>>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
        let analytics = self
            .client
            .begin_scan_analytics(&self.schema, projection, limit);

        let projected_schema = match projection {
            Some(indices) => Arc::new(self.schema.project(indices)?),
            None => Arc::clone(&self.schema),
        };
        let projected_columns = projection.map(|indices| {
            let mut seen = HashSet::new();
            indices
                .iter()
                .filter_map(|index| {
                    let name = self.schema.field(*index).name();
                    seen.insert(name).then(|| name.clone())
                })
                .collect::<Arc<[_]>>()
        });
        let params = ScanParams {
            projected_columns,
            filters: Arc::from(filters),
            limit,
        };

        StreamingTableExec::try_new(
            Arc::clone(&projected_schema),
            vec![Arc::new(GrpcStreamPartitionStream::new(
                &projected_schema,
                self.client.clone(),
                analytics,
                params,
            ))],
            None,
            Vec::default(),
            false,
            None,
        )
        .map(|e| Arc::new(e) as Arc<dyn ExecutionPlan>)
        .map(|exec| {
            Arc::new(SizedCoalesceBatchesExec::new(
                exec,
                CoalescerOptions {
                    target_batch_rows: crate::dataframe_query_common::DEFAULT_BATCH_ROWS,
                    target_batch_bytes: crate::dataframe_query_common::DEFAULT_BATCH_BYTES,
                    max_rows: limit,
                },
            )) as Arc<dyn ExecutionPlan>
        })
    }

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
        self.client.supports_filters_pushdown(filters)
    }

    async fn insert_into(
        &self,
        state: &dyn Session,
        input: Arc<dyn ExecutionPlan>,
        insert_op: InsertOp,
    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
        self.client.insert_into(state, input, insert_op).await
    }
}

#[derive(Debug)]
pub struct GrpcStreamPartitionStream<T: GrpcStreamToTable> {
    schema: SchemaRef,
    client: T,
    analytics: Option<PendingTableQueryAnalytics>,
    params: ScanParams,
}

impl<T: GrpcStreamToTable> GrpcStreamPartitionStream<T> {
    fn new(
        schema: &SchemaRef,
        client: T,
        analytics: Option<PendingTableQueryAnalytics>,
        params: ScanParams,
    ) -> Self {
        Self {
            schema: Arc::clone(schema),
            client,
            analytics,
            params,
        }
    }
}

impl<T> PartitionStream for GrpcStreamPartitionStream<T>
where
    T: GrpcStreamToTable + Send + 'static,
    T::GrpcStreamData: Send + 'static,
{
    fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
        Box::pin(GrpcStream::execute(
            &self.schema,
            self.client.clone(),
            self.analytics.clone(),
            self.params.clone(),
        ))
    }
}

pub struct GrpcStream {
    schema: SchemaRef,
    adapted_stream: Pin<Box<dyn Stream<Item = datafusion::common::Result<RecordBatch>> + Send>>,
}

impl GrpcStream {
    fn execute<T>(
        schema: &SchemaRef,
        mut client: T,
        analytics: Option<PendingTableQueryAnalytics>,
        params: ScanParams,
    ) -> Self
    where
        T::GrpcStreamData: Send + 'static,
        T: GrpcStreamToTable + Send + 'static,
    {
        let expected_schema = Arc::clone(schema);
        let adapted_stream = Box::pin(async_stream::try_stream! {
            // Retry only the *opening* of the stream, never the consume loop below.
            //
            // Admission-control endpoints (`ScanSegmentTable`, `QueryDataset`, …) reject with
            // `ResourceExhausted` fail-fast at the handler entry, before any response byte is
            // produced — so a rejected attempt provably did no server-side work and is safe to
            // retry. Each attempt clones the (read-only/idempotent) provider so the request is
            // rebuilt fresh.
            let mut stream = re_redap_client::with_retry_resource_exhausted("grpc_stream_open", || {
                    let mut client = client.clone();
                    let params = &params;
                    async move { client.send_streaming_request(params).await }
                })
                .await
                .map_err(|err| {
                    if let Some(analytics) = analytics.as_ref() {
                        analytics.record_error(QueryErrorKind::GrpcFetch);
                    }
                    err.into_df_error()
                })?;

            let trace_id = stream.trace_id();
            if let (Some(analytics), Some(trace_id)) = (analytics.as_ref(), trace_id) {
                analytics.record_trace_id(trace_id);
            }

            while let Some(msg) = stream.next().await {
                let msg = msg.map_err(|err| {
                        if let Some(analytics) = analytics.as_ref() {
                            analytics.record_error(QueryErrorKind::GrpcFetch);
                        }
                        err.into_df_error()
                    })?;
                if let Some(analytics) = analytics.as_ref() {
                    analytics.record_first_response();
                }
                let processed = client
                    .process_response(msg, &params)
                    .and_then(|batch| normalize_batch(batch, &expected_schema))
                    .map_err(|err| {
                        if let Some(analytics) = analytics.as_ref() {
                            analytics.record_error(QueryErrorKind::Decode);
                        }
                        err.with_trace_id(trace_id).into_df_error()
                    })?;
                if let Some(analytics) = analytics.as_ref() {
                    analytics.record_first_batch();
                    let num_rows = processed.num_rows() as u64;
                    let num_bytes: u64 = processed
                        .columns()
                        .iter()
                        .map(|c| c.get_array_memory_size() as u64)
                        .sum();
                    analytics.record_batch(num_rows, num_bytes);
                }
                yield processed;
            }
        });

        Self {
            schema: Arc::clone(schema),
            adapted_stream,
        }
    }
}

fn normalize_batch(batch: RecordBatch, expected_schema: &SchemaRef) -> ApiResult<RecordBatch> {
    if batch.schema_ref() == expected_schema {
        return Ok(batch);
    }

    let actual_schema = batch.schema();
    let columns = expected_schema
        .fields()
        .iter()
        .map(|expected_field| {
            let mut matching_columns = actual_schema
                .fields()
                .iter()
                .enumerate()
                .filter(|(_, field)| field.name() == expected_field.name());
            let (index, actual_field) = matching_columns.next().ok_or_else(|| {
                re_redap_client::ApiError::deserialization(
                    None,
                    format!(
                        "response batch is missing projected column '{}'",
                        expected_field.name()
                    ),
                )
            })?;
            if matching_columns.next().is_some() {
                return Err(re_redap_client::ApiError::deserialization(
                    None,
                    format!(
                        "response batch contains projected column '{}' more than once",
                        expected_field.name()
                    ),
                ));
            }
            if actual_field.data_type() != expected_field.data_type() {
                return Err(re_redap_client::ApiError::deserialization(
                    None,
                    format!(
                        "response column '{}' has datatype {}, expected {}",
                        expected_field.name(),
                        actual_field.data_type(),
                        expected_field.data_type()
                    ),
                ));
            }
            Ok(Arc::clone(batch.column(index)))
        })
        .collect::<ApiResult<Vec<_>>>()?;

    RecordBatch::try_new_with_options(
        Arc::clone(expected_schema),
        columns,
        &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
    )
    .map_err(|err| {
        re_redap_client::ApiError::deserialization_with_source(
            None,
            err,
            "failed to normalize projected response batch",
        )
    })
}

impl RecordBatchStream for GrpcStream {
    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }
}

impl Stream for GrpcStream {
    type Item = DataFusionResult<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        self.adapted_stream.poll_next_unpin(cx)
    }
}

#[cfg(test)]
mod table_query_pipeline_tests {
    //! End-to-end tests that drive [`GrpcStream::execute`] with a deterministic
    //! fake [`GrpcStreamToTable`] impl and assert on the analytics state
    //! recorded into the resulting OTLP span.

    use std::collections::HashMap;
    use std::sync::Arc;

    use arrow::array::{RecordBatchOptions, StringArray, UInt32Array};
    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    use async_trait::async_trait;
    use futures_util::StreamExt as _;
    use parking_lot::Mutex;
    use re_redap_client::{ApiError, ApiResponseStream, ApiResult};
    use re_uri::Origin;

    use crate::analytics::{TableQueryInfo, build_table_query_span};
    use crate::{ConnectionAnalytics, PendingTableQueryAnalytics, TableKind, TableQueryCaller};

    use super::*;

    /// Test-only [`GrpcStreamToTable`] impl with deterministic, configurable
    /// behavior. Owns a queue of items the stream will yield, plus knobs to
    /// fail at each stage of the pipeline.
    #[derive(Debug, Clone)]
    struct FakeProvider {
        /// Items the stream yields, consumed on the first `send_streaming_request` call.
        /// `Ok(v)` ⇒ yields a message that decodes to a single-row batch with `v`.
        /// `Err(_)` ⇒ yields a stream-level error (simulates a mid-stream gRPC failure).
        items: Arc<Mutex<Option<Vec<ApiResult<u32>>>>>,
        fail_send_request: bool,
        fail_decode: bool,

        /// Number of `send_streaming_request` opens to reject with `ResourceExhausted` before
        /// succeeding. `Arc<AtomicUsize>` so the count is shared across the per-attempt clones the
        /// retry wrapper makes.
        resource_exhausted_remaining: Arc<std::sync::atomic::AtomicUsize>,
        trace_id: Option<opentelemetry::TraceId>,
    }

    impl FakeProvider {
        fn new(items: Vec<ApiResult<u32>>) -> Self {
            Self {
                items: Arc::new(Mutex::new(Some(items))),
                fail_send_request: false,
                fail_decode: false,
                resource_exhausted_remaining: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
                trace_id: None,
            }
        }

        fn with_trace_id(mut self, trace_id: opentelemetry::TraceId) -> Self {
            self.trace_id = Some(trace_id);
            self
        }

        fn fail_send_request() -> Self {
            Self {
                items: Arc::new(Mutex::new(Some(vec![]))),
                fail_send_request: true,
                fail_decode: false,
                resource_exhausted_remaining: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
                trace_id: None,
            }
        }

        fn fail_decode(items: Vec<u32>) -> Self {
            Self {
                items: Arc::new(Mutex::new(Some(
                    items.into_iter().map(Ok).collect::<Vec<_>>(),
                ))),
                fail_send_request: false,
                fail_decode: true,
                resource_exhausted_remaining: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
                trace_id: None,
            }
        }

        /// Reject the first `n` opens with `ResourceExhausted`, then stream `items`.
        fn resource_exhausted_then(n: usize, items: Vec<u32>) -> Self {
            Self {
                items: Arc::new(Mutex::new(Some(
                    items.into_iter().map(Ok).collect::<Vec<_>>(),
                ))),
                fail_send_request: false,
                fail_decode: false,
                resource_exhausted_remaining: Arc::new(std::sync::atomic::AtomicUsize::new(n)),
                trace_id: None,
            }
        }
    }

    fn fake_schema() -> SchemaRef {
        Arc::new(Schema::new_with_metadata(
            vec![Field::new("v", DataType::UInt32, false)],
            HashMap::default(),
        ))
    }

    fn projection_test_batch() -> RecordBatch {
        RecordBatch::try_new_with_options(
            Arc::new(Schema::new_with_metadata(
                vec![
                    Field::new("number", DataType::UInt32, false),
                    Field::new("label", DataType::Utf8, false),
                ],
                HashMap::default(),
            )),
            vec![
                Arc::new(UInt32Array::from(vec![1, 2])),
                Arc::new(StringArray::from(vec!["one", "two"])),
            ],
            &RecordBatchOptions::new().with_row_count(Some(2)),
        )
        .unwrap()
    }

    #[test]
    fn normalize_batch_projects_reorders_and_uses_expected_schema() {
        let expected_schema = Arc::new(Schema::new_with_metadata(
            vec![
                Field::new("label", DataType::Utf8, false),
                Field::new("number", DataType::UInt32, false),
            ],
            HashMap::from([("source".to_owned(), "catalog".to_owned())]),
        ));

        let normalized = normalize_batch(projection_test_batch(), &expected_schema).unwrap();

        assert_eq!(normalized.schema_ref(), &expected_schema);
        assert_eq!(normalized.num_rows(), 2);
        assert_eq!(
            normalized
                .column(0)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap()
                .value(1),
            "two"
        );
        assert_eq!(
            normalized
                .column(1)
                .as_any()
                .downcast_ref::<UInt32Array>()
                .unwrap()
                .value(0),
            1
        );
    }

    #[test]
    fn normalize_batch_preserves_row_count_for_zero_column_projection() {
        let expected_schema = Arc::new(Schema::empty());

        let normalized = normalize_batch(projection_test_batch(), &expected_schema).unwrap();

        assert_eq!(normalized.num_columns(), 0);
        assert_eq!(normalized.num_rows(), 2);
    }

    #[test]
    fn normalize_batch_rejects_missing_projected_column() {
        let expected_schema = Arc::new(Schema::new_with_metadata(
            vec![Field::new("missing", DataType::Utf8, false)],
            HashMap::default(),
        ));

        let err = normalize_batch(projection_test_batch(), &expected_schema).unwrap_err();

        assert!(
            err.to_string()
                .contains("missing projected column 'missing'")
        );
    }

    #[test]
    fn normalize_batch_rejects_incompatible_datatype() {
        let expected_schema = Arc::new(Schema::new_with_metadata(
            vec![Field::new("number", DataType::Utf8, false)],
            HashMap::default(),
        ));

        let err = normalize_batch(projection_test_batch(), &expected_schema).unwrap_err();

        assert!(
            err.to_string()
                .contains("has datatype UInt32, expected Utf8")
        );
    }

    #[derive(Debug, Clone)]
    struct ProjectionProvider {
        requested_columns: Arc<Mutex<Vec<Option<Vec<String>>>>>,
    }

    #[async_trait]
    impl GrpcStreamToTable for ProjectionProvider {
        type GrpcStreamData = RecordBatch;

        async fn fetch_schema(&mut self) -> ApiResult<SchemaRef> {
            Ok(projection_test_batch().schema())
        }

        fn process_response(
            &mut self,
            response: RecordBatch,
            _params: &ScanParams,
        ) -> ApiResult<RecordBatch> {
            Ok(response)
        }

        async fn send_streaming_request(
            &mut self,
            params: &ScanParams,
        ) -> ApiResult<ApiResponseStream<Self::GrpcStreamData>> {
            self.requested_columns
                .lock()
                .push(params.projected_columns.as_deref().map(<[String]>::to_vec));
            Ok(ApiResponseStream::new(
                futures_util::stream::once(async { Ok(projection_test_batch()) }),
                None,
            ))
        }
    }

    #[tokio::test]
    async fn scan_pushes_projection_and_normalizes_full_response() {
        let requested_columns = Arc::new(Mutex::new(Vec::new()));
        let provider = GrpcStreamProvider::prepare(ProjectionProvider {
            requested_columns: Arc::clone(&requested_columns),
        })
        .await
        .unwrap();
        let context = datafusion::execution::context::SessionContext::new();
        let projection = vec![1];

        let plan = provider
            .scan(&context.state(), Some(&projection), &[], None)
            .await
            .unwrap();
        let batches = datafusion::physical_plan::collect(plan, context.task_ctx())
            .await
            .unwrap();

        assert_eq!(
            *requested_columns.lock(),
            vec![Some(vec!["label".to_owned()])]
        );
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].schema().field(0).name(), "label");
        assert_eq!(batches[0].num_columns(), 1);
        assert_eq!(batches[0].num_rows(), 2);
    }

    #[tokio::test]
    async fn scan_preserves_rows_for_zero_column_projection() {
        let requested_columns = Arc::new(Mutex::new(Vec::new()));
        let provider = GrpcStreamProvider::prepare(ProjectionProvider {
            requested_columns: Arc::clone(&requested_columns),
        })
        .await
        .unwrap();
        let context = datafusion::execution::context::SessionContext::new();
        let projection = Vec::new();

        let plan = provider
            .scan(&context.state(), Some(&projection), &[], None)
            .await
            .unwrap();
        let batches = datafusion::physical_plan::collect(plan, context.task_ctx())
            .await
            .unwrap();

        assert_eq!(*requested_columns.lock(), vec![Some(Vec::new())]);
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].num_columns(), 0);
        assert_eq!(batches[0].num_rows(), 2);
    }

    #[async_trait]
    impl GrpcStreamToTable for FakeProvider {
        type GrpcStreamData = u32;

        async fn fetch_schema(&mut self) -> ApiResult<SchemaRef> {
            Ok(fake_schema())
        }

        fn process_response(
            &mut self,
            response: u32,
            _params: &ScanParams,
        ) -> ApiResult<RecordBatch> {
            if self.fail_decode {
                return Err(ApiError::deserialization(None, "fake decode error"));
            }
            let arr = UInt32Array::from(vec![response]);
            RecordBatch::try_new_with_options(
                fake_schema(),
                vec![Arc::new(arr)],
                &RecordBatchOptions::new().with_row_count(Some(1)),
            )
            .map_err(|err| ApiError::internal_with_source(None, err, "build batch"))
        }

        async fn send_streaming_request(
            &mut self,
            _params: &ScanParams,
        ) -> ApiResult<ApiResponseStream<Self::GrpcStreamData>> {
            // Simulate fail-fast admission control: reject the first `n` opens with
            // ResourceExhausted (shared across the retry wrapper's per-attempt clones). The branch
            // returns before touching `items`, so the eventual successful open still streams them.
            if self
                .resource_exhausted_remaining
                .load(std::sync::atomic::Ordering::SeqCst)
                > 0
            {
                self.resource_exhausted_remaining
                    .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
                return Err(ApiError::tonic(
                    tonic::Status::resource_exhausted("busy"),
                    "fake",
                ));
            }
            if self.fail_send_request {
                return Err(ApiError::deserialization(None, "fake send error"));
            }
            let items = self.items.lock().take().unwrap_or_default();
            let stream = futures_util::stream::iter(items);
            Ok(ApiResponseStream::new(stream, self.trace_id))
        }
    }

    fn make_pending() -> PendingTableQueryAnalytics {
        let origin: Origin = "rerun+http://localhost:51234".parse().unwrap();
        let analytics = ConnectionAnalytics::disabled_for_test(origin);
        analytics.begin_table_query(
            TableQueryInfo {
                table_id: "tbl-pipeline".to_owned(),
                table_kind: TableKind::Lance,
                caller: TableQueryCaller::CatalogResolver,
                schema_total_columns: 1,
                projected_columns: 1,
                has_limit: false,
                limit_value: None,
                time_range: web_time::SystemTime::now()..web_time::SystemTime::now(),
                filters_total: 0,
                filters_signatures: String::new(),
            },
            web_time::Instant::now(),
        )
    }

    fn find_int(span: &opentelemetry_proto::tonic::trace::v1::Span, key: &str) -> Option<i64> {
        use opentelemetry_proto::tonic::common::v1::any_value::Value;
        span.attributes
            .iter()
            .find(|kv| kv.key == key)
            .and_then(|kv| match kv.value.as_ref()?.value.as_ref()? {
                Value::IntValue(i) => Some(*i),
                _ => None,
            })
    }

    fn find_string<'a>(
        span: &'a opentelemetry_proto::tonic::trace::v1::Span,
        key: &str,
    ) -> Option<&'a str> {
        use opentelemetry_proto::tonic::common::v1::any_value::Value;
        span.attributes
            .iter()
            .find(|kv| kv.key == key)
            .and_then(|kv| match kv.value.as_ref()?.value.as_ref()? {
                Value::StringValue(s) => Some(s.as_str()),
                _ => None,
            })
    }

    fn find_bool(span: &opentelemetry_proto::tonic::trace::v1::Span, key: &str) -> Option<bool> {
        use opentelemetry_proto::tonic::common::v1::any_value::Value;
        span.attributes
            .iter()
            .find(|kv| kv.key == key)
            .and_then(|kv| match kv.value.as_ref()?.value.as_ref()? {
                Value::BoolValue(b) => Some(*b),
                _ => None,
            })
    }

    /// Drive a [`GrpcStream`] to completion (or first error) and return the
    /// collected results.
    async fn drain(stream: GrpcStream) -> Vec<DataFusionResult<RecordBatch>> {
        let mut stream = stream;
        let mut out = Vec::new();
        while let Some(item) = stream.next().await {
            let is_err = item.is_err();
            out.push(item);
            // Stop after the first error — `try_stream!` ends the stream there too.
            if is_err {
                break;
            }
        }
        out
    }

    #[tokio::test]
    async fn pipeline_records_per_batch_stats_and_first_response() {
        let provider = FakeProvider::new(vec![Ok(1), Ok(2), Ok(3)]);
        let pending = make_pending();
        let stream = GrpcStream::execute(
            &fake_schema(),
            provider,
            Some(pending.clone()),
            ScanParams::default(),
        );

        let items = drain(stream).await;
        assert_eq!(items.len(), 3);
        assert!(items.iter().all(|r| r.is_ok()));

        let span = pending.build_span_for_test();
        assert_eq!(find_int(&span, "fetch_grpc_requests"), Some(3));
        assert_eq!(find_int(&span, "num_record_batches"), Some(3));
        assert_eq!(find_int(&span, "rows_returned"), Some(3));
        assert!(
            find_int(&span, "bytes_returned").unwrap() > 0,
            "bytes should reflect arrow array size"
        );
        // First response/batch hooks fire.
        assert!(find_int(&span, "time_to_first_response_us").is_some());
        assert!(find_int(&span, "time_to_first_batch_us").is_some());
        assert_eq!(find_bool(&span, "is_success"), Some(true));
    }

    #[tokio::test]
    async fn pipeline_records_grpc_fetch_error_when_send_request_fails() {
        let provider = FakeProvider::fail_send_request();
        let pending = make_pending();
        let stream = GrpcStream::execute(
            &fake_schema(),
            provider,
            Some(pending.clone()),
            ScanParams::default(),
        );

        let items = drain(stream).await;
        assert_eq!(items.len(), 1);
        assert!(items[0].is_err());

        let span = pending.build_span_for_test();
        assert_eq!(find_bool(&span, "is_success"), Some(false));
        assert_eq!(find_string(&span, "error_kind"), Some("grpc_fetch"));
        // No batches were produced.
        assert_eq!(find_int(&span, "num_record_batches"), Some(0));
        assert_eq!(find_int(&span, "rows_returned"), Some(0));
        // First-response was never reached.
        assert!(find_int(&span, "time_to_first_response_us").is_none());
    }

    #[tokio::test]
    async fn pipeline_records_grpc_fetch_error_on_stream_item_error() {
        let provider = FakeProvider::new(vec![
            Ok(1),
            Err(ApiError::deserialization(None, "fake mid-stream err")),
        ]);
        let pending = make_pending();
        let stream = GrpcStream::execute(
            &fake_schema(),
            provider,
            Some(pending.clone()),
            ScanParams::default(),
        );

        let items = drain(stream).await;
        // First batch decoded successfully, second iteration surfaces the error.
        assert_eq!(items.len(), 2);
        assert!(items[0].is_ok());
        assert!(items[1].is_err());

        let span = pending.build_span_for_test();
        assert_eq!(find_bool(&span, "is_success"), Some(false));
        assert_eq!(find_string(&span, "error_kind"), Some("grpc_fetch"));
        // The successful batch before the error is still counted.
        assert_eq!(find_int(&span, "num_record_batches"), Some(1));
        assert_eq!(find_int(&span, "rows_returned"), Some(1));
    }

    #[tokio::test]
    async fn pipeline_records_decode_error() {
        let provider = FakeProvider::fail_decode(vec![1, 2]);
        let pending = make_pending();
        let stream = GrpcStream::execute(
            &fake_schema(),
            provider,
            Some(pending.clone()),
            ScanParams::default(),
        );

        let items = drain(stream).await;
        assert_eq!(items.len(), 1);
        assert!(items[0].is_err());

        let span = pending.build_span_for_test();
        assert_eq!(find_bool(&span, "is_success"), Some(false));
        assert_eq!(find_string(&span, "error_kind"), Some("decode"));
        // First gRPC message arrived (so first_response was set), but no batch
        // was successfully decoded.
        assert!(find_int(&span, "time_to_first_response_us").is_some());
        assert_eq!(find_int(&span, "num_record_batches"), Some(0));
    }

    #[tokio::test]
    async fn pipeline_propagates_trace_id_into_span() {
        let trace_id = opentelemetry::TraceId::from_bytes([9u8; 16]);
        let provider = FakeProvider::new(vec![Ok(1)]).with_trace_id(trace_id);
        let pending = make_pending();
        let stream = GrpcStream::execute(
            &fake_schema(),
            provider,
            Some(pending.clone()),
            ScanParams::default(),
        );

        let _ = drain(stream).await;

        let span = pending.build_span_for_test();
        assert_eq!(span.trace_id, trace_id.to_bytes());
        assert_eq!(span.span_id.len(), 8);
        assert!(span.span_id.iter().any(|byte| *byte != 0));
    }

    #[tokio::test]
    async fn pipeline_runs_without_analytics_attached() {
        // No PendingTableQueryAnalytics — recording paths are skipped, but the
        // stream still produces correct output. Smoke test for the
        // `if let Some(analytics) = analytics.as_ref()` branches.
        let provider = FakeProvider::new(vec![Ok(1), Ok(2)]);
        let stream = GrpcStream::execute(&fake_schema(), provider, None, ScanParams::default());

        let items = drain(stream).await;
        assert_eq!(items.len(), 2);
        assert!(items.iter().all(|r| r.is_ok()));
    }

    #[tokio::test]
    async fn pipeline_retries_resource_exhausted_then_succeeds() {
        // The first two stream-opens are rejected with `ResourceExhausted`; the retry wrapper in
        // `GrpcStream::execute` should ride them out and then stream all three rows.
        let provider = FakeProvider::resource_exhausted_then(2, vec![1, 2, 3]);
        let stream = GrpcStream::execute(&fake_schema(), provider, None, ScanParams::default());

        let items = drain(stream).await;
        assert_eq!(items.len(), 3);
        assert!(items.iter().all(|r| r.is_ok()));
    }

    #[test]
    fn begin_scan_analytics_default_returns_none() {
        // The default trait impl on `GrpcStreamToTable` returns None — only
        // `TableEntryTableProvider` overrides it. Non-table providers should
        // remain analytics-free.
        let provider = FakeProvider::new(vec![]);
        let schema = fake_schema();
        let result = provider.begin_scan_analytics(&schema, None, None);
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn build_table_query_span_called_via_pending_matches_pure_builder() {
        // Sanity check: PendingTableQueryAnalytics::build_span_for_test
        // produces equivalent output to calling `build_table_query_span`
        // directly with the same inputs. Pins the wiring between the two.
        let pending = make_pending();
        pending.record_batch(10, 100);
        pending.record_batch(20, 200);
        pending.record_first_response();
        pending.record_first_batch();

        let span_via_pending = pending.build_span_for_test();
        // Required keys present in both forms.
        let direct = build_table_query_span(
            &TableQueryInfo {
                table_id: "tbl-pipeline".to_owned(),
                table_kind: TableKind::Lance,
                caller: TableQueryCaller::CatalogResolver,
                schema_total_columns: 1,
                projected_columns: 1,
                has_limit: false,
                limit_value: None,
                time_range: web_time::SystemTime::now()..web_time::SystemTime::now(),
                filters_total: 0,
                filters_signatures: String::new(),
            },
            crate::analytics::TableScanStatsSnapshot {
                grpc_requests: 2,
                batches: 2,
                rows_returned: 30,
                bytes_returned: 300,
            },
            web_time::SystemTime::now()..web_time::SystemTime::now(),
            std::time::Duration::ZERO,
            None,
            None,
            None,
            None,
        );

        // Same span name and required-attr counts, regardless of the
        // (slightly different) timing values the wrappers compute.
        assert_eq!(span_via_pending.name, direct.name);
        assert_eq!(find_int(&span_via_pending, "rows_returned"), Some(30));
        assert_eq!(find_int(&span_via_pending, "num_record_batches"), Some(2));
    }
}