/// creates a new PyLogicalPlan

pub fn new(plan: LogicalPlan) -> Self {

Self {

plan: Arc::new(plan),

}

}

pub fn plan(&self) -> Arc<LogicalPlan> {

self.plan.clone()

}

/// Return the specific logical operator

pub fn to_variant(&self, py: Python) -> PyResult<PyObject> {

match self.plan.as_ref() {

LogicalPlan::Aggregate(plan) => PyAggregate::from(plan.clone()).to_variant(py),

LogicalPlan::Analyze(plan) => PyAnalyze::from(plan.clone()).to_variant(py),

LogicalPlan::Distinct(plan) => PyDistinct::from(plan.clone()).to_variant(py),

LogicalPlan::EmptyRelation(plan) => PyEmptyRelation::from(plan.clone()).to_variant(py),

LogicalPlan::Explain(plan) => PyExplain::from(plan.clone()).to_variant(py),

LogicalPlan::Extension(plan) => PyExtension::from(plan.clone()).to_variant(py),

LogicalPlan::Filter(plan) => PyFilter::from(plan.clone()).to_variant(py),

LogicalPlan::Join(plan) => PyJoin::from(plan.clone()).to_variant(py),

LogicalPlan::Limit(plan) => PyLimit::from(plan.clone()).to_variant(py),

LogicalPlan::Projection(plan) => PyProjection::from(plan.clone()).to_variant(py),

LogicalPlan::Sort(plan) => PySort::from(plan.clone()).to_variant(py),

LogicalPlan::TableScan(plan) => PyTableScan::from(plan.clone()).to_variant(py),

LogicalPlan::Subquery(plan) => PySubquery::from(plan.clone()).to_variant(py),

LogicalPlan::SubqueryAlias(plan) => PySubqueryAlias::from(plan.clone()).to_variant(py),

LogicalPlan::Unnest(plan) => PyUnnest::from(plan.clone()).to_variant(py),

LogicalPlan::Window(plan) => PyWindowExpr::from(plan.clone()).to_variant(py),

LogicalPlan::Repartition(_)

| LogicalPlan::Union(_)

| LogicalPlan::Statement(_)

| LogicalPlan::Values(_)

| LogicalPlan::Dml(_)

| LogicalPlan::Ddl(_)

| LogicalPlan::Copy(_)

| LogicalPlan::DescribeTable(_)

| LogicalPlan::RecursiveQuery(_) => Err(py_unsupported_variant_err(format!(

"Conversion of variant not implemented: {:?}",

self.plan

))),

}

}

/// Get the inputs to this plan

fn inputs(&self) -> Vec<PyLogicalPlan> {

let mut inputs = vec![];

for input in self.plan.inputs() {

inputs.push(input.to_owned().into());

}

inputs

}

fn __repr__(&self) -> PyResult<String> {

Ok(format!("{:?}", self.plan))

}

fn display(&self) -> String {

format!("{}", self.plan.display())

}

fn display_indent(&self) -> String {

format!("{}", self.plan.display_indent())

}

fn display_indent_schema(&self) -> String {

format!("{}", self.plan.display_indent_schema())

}

fn display_graphviz(&self) -> String {

format!("{}", self.plan.display_graphviz())

}

pub fn to_proto<'py>(&'py self, py: Python<'py>) -> PyDataFusionResult<Bound<'py, PyBytes>> {

let codec = DefaultLogicalExtensionCodec {};

let proto =

datafusion_proto::protobuf::LogicalPlanNode::try_from_logical_plan(&self.plan, &codec)?;

let bytes = proto.encode_to_vec();

Ok(PyBytes::new_bound(py, &bytes))

}

#[staticmethod]

pub fn from_proto(

ctx: PySessionContext,

proto_msg: Bound<'_, PyBytes>,

) -> PyDataFusionResult<Self> {

let bytes: &[u8] = proto_msg.extract()?;

let proto_plan =

datafusion_proto::protobuf::LogicalPlanNode::decode(bytes).map_err(|e| {

PyRuntimeError::new_err(format!(

"Unable to decode logical node from serialized bytes: {}",

e

))

})?;

let codec = DefaultLogicalExtensionCodec {};

let plan = proto_plan.try_into_logical_plan(&ctx.ctx, &codec)?;

Ok(Self::new(plan))

}

fn from(logical_plan: PyLogicalPlan) -> LogicalPlan {

logical_plan.plan.as_ref().clone()

}

fn from(logical_plan: LogicalPlan) -> PyLogicalPlan {

PyLogicalPlan {

plan: Arc::new(logical_plan),

}

}