Add missing imports for Python bindings

crates/modelardb_bulkloader/src/main.rs

@@ -218,19 +218,18 @@ async fn import_time_series_table(
         // Write the current batch if it uses more than half the memory so concatenation is
         // possible. The amount of available memory is reduced by 20% for other variables.
         system.refresh_memory();
-        if current_batch_size > (system.available_memory() as usize / 10 * 8) {
-            if let Err(write_error) = import_and_clear_time_series_table_batch(
+        if current_batch_size > (system.available_memory() as usize / 10 * 8)
+            && let Err(write_error) = import_and_clear_time_series_table_batch(
                 data_folder,
                 &mut delta_table_writer,
                 time_series_table_metadata,
                 &mut current_batch,
                 &mut current_batch_size,
             )
             .await
-            {
-                delta_table_writer.rollback().await?;
-                return Err(write_error);
-            }
+        {
+            delta_table_writer.rollback().await?;
+            return Err(write_error);
         }
     }
 

crates/modelardb_embedded/bindings/python/modelardb/__init__.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .operations import Aggregate, Operations
+from .operations import Aggregate, Operations, open_memory, open_local, open_s3, open_azure, connect
 from .node import Server, Manager
 from .error_bound import AbsoluteErrorBound, RelativeErrorBound
 from .table import NormalTable, TimeSeriesTable

crates/modelardb_server/tests/integration_test.rs

@@ -17,14 +17,14 @@
 
 use std::collections::HashMap;
 use std::error::Error;
-use std::iter;
 use std::ops::Range;
 use std::process::Stdio;
 use std::str;
 use std::string::String;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU16, Ordering};
 use std::time::Duration;
+use std::{iter, slice};
 
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::{Action, Criteria, FlightData, FlightDescriptor, PutResult, Ticket, utils};
@@ -626,7 +626,7 @@ async fn test_can_truncate_normal_table() {
 
     ingest_time_series_and_flush_data(
         &mut test_context,
-        &[time_series.clone()],
+        slice::from_ref(&time_series),
         TableType::NormalTable,
     )
     .await;
@@ -649,7 +649,7 @@ async fn test_can_truncate_time_series_table() {
 
     ingest_time_series_and_flush_data(
         &mut test_context,
-        &[time_series.clone()],
+        slice::from_ref(&time_series),
         TableType::TimeSeriesTable,
     )
     .await;
@@ -736,7 +736,7 @@ async fn test_do_put_can_ingest_time_series_with_tags() {
 
     ingest_time_series_and_flush_data(
         &mut test_context,
-        &[time_series.clone()],
+        slice::from_ref(&time_series),
         TableType::TimeSeriesTable,
     )
     .await;
@@ -787,7 +787,7 @@ async fn test_do_put_can_ingest_time_series_without_tags() {
 
     ingest_time_series_and_flush_data(
         &mut test_context,
-        &[time_series.clone()],
+        slice::from_ref(&time_series),
         TableType::TimeSeriesTableNoTag,
     )
     .await;
@@ -838,7 +838,7 @@ async fn test_do_put_can_ingest_time_series_with_generated_field() {
 
     ingest_time_series_and_flush_data(
         &mut test_context,
-        &[time_series.clone()],
+        slice::from_ref(&time_series),
         TableType::TimeSeriesTableAsField,
     )
     .await;

crates/modelardb_storage/src/optimizer/model_simple_aggregates.rs

@@ -208,43 +208,42 @@ fn rewrite_aggregates_to_use_segments(
     // The rule tries to match the subtree of execution_plan so execution_plan can be updated.
     let execution_plan_children = execution_plan.children();
 
-    if execution_plan_children.len() == 1 {
-        if let Some(aggregate_exec) = execution_plan_children[0]
+    if execution_plan_children.len() == 1
+        && let Some(aggregate_exec) = execution_plan_children[0]
             .as_any()
             .downcast_ref::<AggregateExec>()
+    {
+        // Currently, only aggregates on one FIELD column without predicates are supported.
+        let aggregate_exec_children = aggregate_exec.children();
+        if aggregate_exec.input_schema().fields.len() == 1
+            && *aggregate_exec.input_schema().field(0).data_type() == ArrowValue::DATA_TYPE
+            && aggregate_exec.filter_expr().iter().all(Option::is_none)
+            && aggregate_exec.group_expr().is_empty()
         {
-            // Currently, only aggregates on one FIELD column without predicates are supported.
-            let aggregate_exec_children = aggregate_exec.children();
-            if aggregate_exec.input_schema().fields.len() == 1
-                && *aggregate_exec.input_schema().field(0).data_type() == ArrowValue::DATA_TYPE
-                && aggregate_exec.filter_expr().iter().all(Option::is_none)
-                && aggregate_exec.group_expr().is_empty()
+            // Remove RepartitionExec if added by Apache DataFusion. Both AggregateExec and
+            // RepartitionExec can only have one child, so it is not necessary to check it.
+            let maybe_repartition_exec = &aggregate_exec_children[0];
+            let aggregate_exec_input = if let Some(repartition_exec) = maybe_repartition_exec
+                .as_any()
+                .downcast_ref::<RepartitionExec>()
             {
-                // Remove RepartitionExec if added by Apache DataFusion. Both AggregateExec and
-                // RepartitionExec can only have one child, so it is not necessary to check it.
-                let maybe_repartition_exec = &aggregate_exec_children[0];
-                let aggregate_exec_input = if let Some(repartition_exec) = maybe_repartition_exec
-                    .as_any()
-                    .downcast_ref::<RepartitionExec>()
-                {
-                    repartition_exec.children()[0].clone()
-                } else {
-                    (*maybe_repartition_exec).clone()
-                };
+                repartition_exec.children()[0].clone()
+            } else {
+                (*maybe_repartition_exec).clone()
+            };
 
-                if let Some(sorted_join_exec) = aggregate_exec_input
-                    .as_any()
-                    .downcast_ref::<SortedJoinExec>()
+            if let Some(sorted_join_exec) = aggregate_exec_input
+                .as_any()
+                .downcast_ref::<SortedJoinExec>()
+            {
+                // Try to create new AggregateExec that compute aggregates directly from segments.
+                if let Ok(input) =
+                    try_new_aggregate_exec(aggregate_exec, sorted_join_exec.children())
                 {
-                    // Try to create new AggregateExec that compute aggregates directly from segments.
-                    if let Ok(input) =
-                        try_new_aggregate_exec(aggregate_exec, sorted_join_exec.children())
-                    {
-                        return Ok(Transformed::yes(
-                            execution_plan.with_new_children(vec![input])?,
-                        ));
-                    };
-                }
+                    return Ok(Transformed::yes(
+                        execution_plan.with_new_children(vec![input])?,
+                    ));
+                };
             }
         }
     }
@@ -285,22 +284,18 @@ fn try_new_aggregate_exec(
 /// Return [`Ok`] if no predicates have been pushed to `grid_exec_child`, otherwise
 /// [`DataFusionError`] is returned.
 fn can_rewrite_aggregate(grid_exec_child: &Arc<dyn ExecutionPlan>) -> DataFusionResult<()> {
-    if let Some(data_source_exec) = grid_exec_child.as_any().downcast_ref::<DataSourceExec>() {
-        if let Some(file_scan_config) = data_source_exec
+    if let Some(data_source_exec) = grid_exec_child.as_any().downcast_ref::<DataSourceExec>()
+        && let Some(file_scan_config) = data_source_exec
             .data_source()
             .as_any()
             .downcast_ref::<FileScanConfig>()
-        {
-            if let Some(parquet_source) = file_scan_config
-                .file_source
-                .as_any()
-                .downcast_ref::<ParquetSource>()
-            {
-                if parquet_source.predicate().is_none() {
-                    return Ok(());
-                }
-            }
-        }
+        && let Some(parquet_source) = file_scan_config
+            .file_source
+            .as_any()
+            .downcast_ref::<ParquetSource>()
+        && parquet_source.predicate().is_none()
+    {
+        return Ok(());
     }
 
     Err(DataFusionError::Plan(

crates/modelardb_storage/src/parser.rs

@@ -182,22 +182,22 @@ impl ModelarDbDialect {
     /// [`false`] is returned. The method does not consume tokens.
     fn next_tokens_are_create_time_series_table(&self, parser: &Parser) -> bool {
         // CREATE.
-        if let Token::Word(word) = parser.peek_nth_token(0).token {
-            if word.keyword == Keyword::CREATE {
-                // TIME.
-                if let Token::Word(word) = parser.peek_nth_token(1).token {
-                    if word.value.to_uppercase() == "TIME" {
-                        // SERIES.
-                        if let Token::Word(word) = parser.peek_nth_token(2).token {
-                            if word.value.to_uppercase() == "SERIES" {
-                                // TABLE.
-                                if let Token::Word(word) = parser.peek_nth_token(3).token {
-                                    if word.keyword == Keyword::TABLE {
-                                        return true;
-                                    }
-                                }
-                            }
-                        }
+        if let Token::Word(word) = parser.peek_nth_token(0).token
+            && word.keyword == Keyword::CREATE
+        {
+            // TIME.
+            if let Token::Word(word) = parser.peek_nth_token(1).token
+                && word.value.to_uppercase() == "TIME"
+            {
+                // SERIES.
+                if let Token::Word(word) = parser.peek_nth_token(2).token
+                    && word.value.to_uppercase() == "SERIES"
+                {
+                    // TABLE.
+                    if let Token::Word(word) = parser.peek_nth_token(3).token
+                        && word.keyword == Keyword::TABLE
+                    {
+                        return true;
                     }
                 }
             }
@@ -298,10 +298,10 @@ impl ModelarDbDialect {
     /// Return [`Ok`] if the next [`Token`] is a [`Token::Word`] with the value `expected`,
     /// otherwise a [`ParserError`] is returned.
     fn expect_word_value(&self, parser: &mut Parser, expected: &str) -> StdResult<(), ParserError> {
-        if let Ok(string) = self.parse_word_value(parser) {
-            if string.to_uppercase() == expected.to_uppercase() {
-                return Ok(());
-            }
+        if let Ok(string) = self.parse_word_value(parser)
+            && string.to_uppercase() == expected.to_uppercase()
+        {
+            return Ok(());
         }
         parser.expected(expected, parser.peek_token())
     }

crates/modelardb_storage/src/query/grid_exec.rs

@@ -21,6 +21,7 @@ use std::any::Any;
 use std::borrow::Cow;
 use std::fmt::{Formatter, Result as FmtResult};
 use std::pin::Pin;
+use std::slice;
 use std::sync::Arc;
 use std::task::{Context as StdTaskContext, Poll};
 
@@ -86,7 +87,7 @@ impl GridExec {
         // assumes the data it receives from all of its inputs uses the same sort order.
         let equivalence_properties = EquivalenceProperties::new_with_orderings(
             schema.clone(),
-            &[query_order_data_point.clone()],
+            slice::from_ref(&query_order_data_point),
         );
 
         let plan_properties = PlanProperties::new(

crates/modelardb_storage/src/query/normal_table.rs

@@ -82,7 +82,7 @@ impl TableProvider for NormalTable {
     }
 
     /// Get the [`LogicalPlan`] of this normal table, if available.
-    fn get_logical_plan(&self) -> Option<Cow<LogicalPlan>> {
+    fn get_logical_plan(&self) -> Option<Cow<'_, LogicalPlan>> {
         self.delta_table.get_logical_plan()
     }