datafusion_python/expr/
create_index.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::{
19    fmt::{self, Display, Formatter},
20    sync::Arc,
21};
22
23use datafusion::logical_expr::CreateIndex;
24use pyo3::{prelude::*, IntoPyObjectExt};
25
26use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
27
28use super::{logical_node::LogicalNode, sort_expr::PySortExpr};
29
30#[pyclass(name = "CreateIndex", module = "datafusion.expr", subclass)]
31#[derive(Clone)]
32pub struct PyCreateIndex {
33    create: CreateIndex,
34}
35
36impl From<PyCreateIndex> for CreateIndex {
37    fn from(create: PyCreateIndex) -> Self {
38        create.create
39    }
40}
41
42impl From<CreateIndex> for PyCreateIndex {
43    fn from(create: CreateIndex) -> PyCreateIndex {
44        PyCreateIndex { create }
45    }
46}
47
48impl Display for PyCreateIndex {
49    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
50        write!(f, "CreateIndex: {:?}", self.create.name)
51    }
52}
53
54#[pymethods]
55impl PyCreateIndex {
56    #[new]
57    #[pyo3(signature = (table, columns, unique, if_not_exists, schema, name=None, using=None))]
58    pub fn new(
59        table: String,
60        columns: Vec<PySortExpr>,
61        unique: bool,
62        if_not_exists: bool,
63        schema: PyDFSchema,
64        name: Option<String>,
65        using: Option<String>,
66    ) -> PyResult<Self> {
67        Ok(PyCreateIndex {
68            create: CreateIndex {
69                name,
70                table: table.into(),
71                using,
72                columns: columns.iter().map(|c| c.clone().into()).collect(),
73                unique,
74                if_not_exists,
75                schema: Arc::new(schema.into()),
76            },
77        })
78    }
79
80    pub fn name(&self) -> Option<String> {
81        self.create.name.clone()
82    }
83
84    pub fn table(&self) -> PyResult<String> {
85        Ok(self.create.table.to_string())
86    }
87
88    pub fn using(&self) -> Option<String> {
89        self.create.using.clone()
90    }
91
92    pub fn columns(&self) -> Vec<PySortExpr> {
93        self.create
94            .columns
95            .iter()
96            .map(|c| c.clone().into())
97            .collect()
98    }
99
100    pub fn unique(&self) -> bool {
101        self.create.unique
102    }
103
104    pub fn if_not_exists(&self) -> bool {
105        self.create.if_not_exists
106    }
107
108    pub fn schema(&self) -> PyDFSchema {
109        (*self.create.schema).clone().into()
110    }
111
112    fn __repr__(&self) -> PyResult<String> {
113        Ok(format!("CreateIndex({})", self))
114    }
115
116    fn __name__(&self) -> PyResult<String> {
117        Ok("CreateIndex".to_string())
118    }
119}
120
121impl LogicalNode for PyCreateIndex {
122    fn inputs(&self) -> Vec<PyLogicalPlan> {
123        vec![]
124    }
125
126    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
127        self.clone().into_bound_py_any(py)
128    }
129}